fix(rust): complete installed callback parity

litellm-rust/crates/ai-gateway/src/routes/messages/service.rs

@@ -54,19 +54,19 @@ pub async fn run(
     let request = MessagesRequest {
         model: provider_model,
         body,
-        options: RequestOptions {
-            api_key: (deployment.litellm_params.api_key.as_deref()).map(|value| value.to_string()),
-            api_base: (deployment.litellm_params.api_base.as_deref())
-                .map(|value| value.to_string()),
-            custom_llm_provider: (custom_llm_provider).map(|value| value.to_string()),
-            extra_headers,
-            timeout: None,
-            ..Default::default()
-        },
+    };
+    let options = RequestOptions {
+        api_key: (deployment.litellm_params.api_key.as_deref()).map(|value| value.to_string()),
+        api_base: (deployment.litellm_params.api_base.as_deref()).map(|value| value.to_string()),
+        custom_llm_provider: (custom_llm_provider).map(|value| value.to_string()),
+        extra_headers,
+        timeout: None,
+        ..Default::default()
     };
     if request.body.get("stream").and_then(Value::as_bool) == Some(true) {
         return messages_stream(
             request,
+            &options,
             &LiteLlmRequestContext {
                 ..Default::default()
             },
@@ -77,6 +77,7 @@ pub async fn run(
 
     let response = messages(
         request,
+        &options,
         &LiteLlmRequestContext {
             ..Default::default()
         },

litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs

@@ -86,7 +86,6 @@ impl ProviderAttemptObserver for ProviderObserver {
 
     async fn error(&mut self, input: &ProviderError) -> Result<(), Self::Error> {
         assert!(input.committed);
-        assert!(!input.will_retry);
         assert!(!input.message.is_empty());
         self.events.lock().unwrap().push("error");
         if self.rejected_callback == Some("error") {

litellm-rust/crates/core/src/provider_callbacks/handler.rs

@@ -143,7 +143,6 @@ where
             Error::Http { status, .. } => Some(*status),
             _ => None,
         },
-        will_retry: false,
         ended_at: epoch_seconds(),
     };
     observer.error(&event).await.map_err(callback_error)
@@ -240,7 +239,6 @@ mod tests {
             assert_eq!(event.attempt, 3);
             assert_eq!(event.trace_id.as_deref(), Some("trace-1"));
             assert!(event.committed);
-            assert!(!event.will_retry);
             self.events.push("error");
             Ok(())
         }

litellm-rust/crates/core/src/provider_callbacks/mod.rs

@@ -58,7 +58,6 @@ pub struct ProviderError {
     pub stage: &'static str,
     pub committed: bool,
     pub status_code: Option<u16>,
-    pub will_retry: bool,
     pub ended_at: f64,
 }
 

litellm-rust/crates/python-bridge/tests/callbacks/mod.rs

@@ -370,7 +370,6 @@ fn provider_error() -> ProviderError {
         stage: "provider_response",
         committed: true,
         status_code: Some(429),
-        will_retry: true,
         ended_at: 2.0,
     }
 }

litellm/ocr/main.py

@@ -30,6 +30,7 @@ from litellm.llms.base_llm.ocr.transformation import (
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.rust_bridge import ocr as rust_ocr_bridge
+from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter
 from litellm.rust_bridge.request import (
     NativeRequestCapabilities,
     NativeRequestOptions,
@@ -240,27 +241,8 @@ def _prepare_rust_ocr_call(
         api_base=prepared_request.api_base,
         litellm_params=prepared_request.litellm_params,
     )
-    resolved_complete_url: Final = provider_config.get_complete_url(
-        api_base=prepared_request.api_base,
-        model=prepared_request.model,
-        optional_params=prepared_request.optional_params,
-        litellm_params=prepared_request.litellm_params,
-    )
     rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
     rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
-    prepared_request.litellm_logging_obj.pre_call(
-        input="OCR document processing",
-        api_key=resolved_api_key,
-        additional_args={
-            "complete_input_dict": {
-                "model": prepared_request.model,
-                "document": prepared_request.document,
-                **rust_optional_params,
-            },
-            "api_base": resolved_complete_url,
-            "headers": resolved_headers,
-        },
-    )
     return PreparedNativeCall(
         request=rust_ocr_bridge.NativeOCRRequest(
             model=prepared_request.model,
@@ -292,6 +274,11 @@ def _prepare_rust_ocr_call(
                 native_response_format=(prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"),
             ),
         ),
+        callback_adapter=ProviderLoggingAdapter(
+            prepared_request.litellm_logging_obj,
+            "OCR document processing",
+            resolved_api_key,
+        ),
     )
 
 

litellm/rust_bridge/callback_adapters.py

@@ -53,7 +53,6 @@ class ProviderError(ProviderEvent):
     stage: str
     committed: bool
     status_code: int | None
-    will_retry: bool
     ended_at: float
 
 

litellm/rust_bridge/callbacks.py

@@ -22,12 +22,11 @@ CallbackDecision: TypeAlias = CallbackUnchanged | CallbackReplace | CallbackReje
 
 
 class ProviderAttemptCallbackHandle(Protocol):
-    """Observe one provider attempt.
+    """Observe the provider operation inside one native call.
 
-    Retries repeat ``pre_call`` and ``error`` with incremented attempt metadata.
-    Only the successful attempt receives ``post_call``. These callbacks observe
-    the provider operation; outer SDK success and failure callbacks remain owned
-    by Python after endpoint dispatch completes.
+    Successful operations receive ``pre_call`` and ``post_call`` once. Failed
+    operations receive ``pre_call`` and ``error`` once. Outer SDK success and
+    failure callbacks remain owned by Python after endpoint dispatch completes.
     """
 
     def pre_call(self, payload: object, /) -> CallbackDecision: ...

litellm/rust_bridge/request.py

@@ -74,6 +74,7 @@ def vertex_options(params: Mapping[str, object]) -> NativeVertexOptions:
         location=location if isinstance(location, str) else None,
     )
 
+
 from typing_extensions import ReadOnly, TypedDict, TypeVar
 
 
@@ -192,6 +193,12 @@ def call_native(
     native: NativeFunction[RequestT, ResultT, CallbackT],
     prepared: PreparedNativeCall[RequestT, CallbackT],
 ) -> ResultT:
+    if prepared.callback_adapter is None:
+        return native(
+            prepared.request,
+            options=prepared.options,
+            context=prepared.context,
+        )
     return native(
         prepared.request,
         options=prepared.options,

tests/test_litellm/ocr/live_callback_smoke.py

@@ -62,5 +62,5 @@ if __name__ == "__main__":
     else:
         native: Final = get_native_bridge()
         assert native is not None
-        with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}):
+        with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}, create=True):
             asyncio.run(smoke(baseline=False))

tests/test_litellm/ocr/sdk_callback_contract.py

@@ -79,7 +79,7 @@ async def exercise(asynchronous: bool, rust: bool, case: str, *, native_expected
             following: Final = tuple(event for event in await follower.wait() if native or event.name != "post")
             names: Final = tuple(event.name for event in events)
             prefix: Final = ("pre", "post") if native and status == 200 and case != "timeout" else ("pre",)
-            assert names[: len(prefix)] == prefix
+            assert names[: len(prefix)] == prefix, (asynchronous, rust, case, native, names, prefix)
             terminal: Final = "success" if successful else "failure"
             expected: Final = (
                 ("async_success",)
@@ -267,7 +267,7 @@ async def verify_foundation() -> None:
     assert_native_unavailable()
     native: Final = get_native_bridge()
     assert native is not None
-    with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}):
+    with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}, create=True):
         await verify_parity()
     assert_native_unavailable()
 

tests/test_litellm/ocr/test_rust_bridge.py

@@ -11,6 +11,7 @@ import pytest
 
 import litellm
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter
 from litellm.rust_bridge import configuration
 from litellm.rust_bridge.request import (
     NativeOCRRequest,
@@ -52,6 +53,7 @@ class RecordingBridge:
 
     def __init__(self) -> None:
         self.calls: list[dict[str, object]] = []
+        self.callback_adapter: object | None = None
 
     def __call__(
         self,
@@ -61,6 +63,7 @@ class RecordingBridge:
         context: NativeRequestContext,
         callback_adapter: object | None = None,
     ) -> dict[str, object]:
+        self.callback_adapter = callback_adapter
         self.calls.append(
             {
                 "model": request.model,
@@ -82,6 +85,7 @@ class RecordingAsyncBridge:
 
     def __init__(self) -> None:
         self.calls: list[dict[str, object]] = []
+        self.callback_adapter: object | None = None
 
     async def __call__(
         self,
@@ -91,6 +95,7 @@ class RecordingAsyncBridge:
         context: NativeRequestContext,
         callback_adapter: object | None = None,
     ) -> dict[str, object]:
+        self.callback_adapter = callback_adapter
         self.calls.append(
             {
                 "model": request.model,
@@ -691,7 +696,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
     assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com"
 
 
-def test_run_rust_ocr_runs_pre_call_logging():
+def test_run_rust_ocr_passes_provider_logging_adapter():
     logging_obj = RecordingLogging()
     bridge = RecordingBridge()
     litellm.rust(True)
@@ -709,17 +714,11 @@ def test_run_rust_ocr_runs_pre_call_logging():
         resolve_api_key=lambda _name: None,
     )
 
-    assert logging_obj.pre_call_kwargs is not None
-    assert logging_obj.pre_call_kwargs["input"] == "OCR document processing"
-    additional_args = logging_obj.pre_call_kwargs["additional_args"]
-    complete_input = additional_args["complete_input_dict"]
-    assert complete_input["document"] == DOCUMENT
-    assert complete_input["include_image_base64"] is True
-    assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr"
-    assert additional_args["headers"] == {
-        "Authorization": "Bearer sk-test",
-        "x-trace-id": "trace-1",
-    }
+    adapter = bridge.callback_adapter
+    assert isinstance(adapter, ProviderLoggingAdapter)
+    assert adapter.logging_obj is logging_obj
+    assert adapter.input == "OCR document processing"
+    assert adapter.api_key == "sk-test"
 
 
 def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge):

tests/test_litellm/rust_bridge/test_callback_adapters.py

@@ -45,7 +45,6 @@ def test_provider_logging_adapter_preserves_provider_lifecycle() -> None:
             stage="provider_response",
             committed=True,
             status_code=429,
-            will_retry=True,
             ended_at=11.0,
         )
     )
@@ -69,7 +68,6 @@ def test_provider_logging_adapter_preserves_provider_lifecycle() -> None:
         "stage": "provider_response",
         "committed": True,
         "status_code": 429,
-        "will_retry": True,
         "ended_at": 11.0,
     }