fix(ocr): retain callback mutations for native calls

litellm/ocr/main.py

@@ -58,8 +58,8 @@ class _PreparedOCRRequest:
 class _PreparedRustOCRCall:
     api_key: str | None
     api_base: str | None
+    body: dict[str, object]
     headers: dict[str, object]
-    optional_params: dict[str, object]
 
 
 _RUST_OCR_PROVIDERS: Final = {
@@ -261,15 +261,16 @@ def _prepare_rust_ocr_call(
     )
     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)
+    body: Final[dict[str, object]] = {
+        "model": prepared_request.model,
+        "document": prepared_request.document,
+        **rust_optional_params,
+    }
     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,
-            },
+            "complete_input_dict": body,
             "api_base": resolved_complete_url,
             "headers": resolved_headers,
         },
@@ -277,11 +278,29 @@ def _prepare_rust_ocr_call(
     return _PreparedRustOCRCall(
         api_key=resolved_api_key,
         api_base=rust_api_base,
+        body=body,
         headers=cast(dict[str, object], resolved_headers),
-        optional_params=rust_optional_params,
     )
 
 
+def _rust_ocr_model(body: Mapping[str, object]) -> str:
+    model: Final = body.get("model")
+    if not isinstance(model, str):
+        raise TypeError("OCR callback produced a non-string model")
+    return model
+
+
+def _rust_ocr_document(body: Mapping[str, object]) -> dict[str, object]:
+    document: Final = body.get("document")
+    if not isinstance(document, dict):
+        raise TypeError("OCR callback produced a non-dict document")
+    return cast(dict[str, object], document)  # cast-ok: the native bridge validates the retained document fields
+
+
+def _rust_ocr_optional_params(body: Mapping[str, object]) -> dict[str, object]:
+    return {name: value for name, value in body.items() if name not in {"model", "document"}}
+
+
 def _map_rust_ocr_error(
     error: Exception,
     prepared_request: _PreparedOCRRequest,
@@ -321,13 +340,13 @@ def _run_rust_ocr(
     )
     try:
         rust_response: Final = rust_ocr_bridge.ocr(
-            model=prepared_request.model,
-            document=prepared_request.document,
+            model=_rust_ocr_model(prepared.body),
+            document=_rust_ocr_document(prepared.body),
             api_key=prepared.api_key,
             api_base=prepared.api_base,
             custom_llm_provider=prepared_request.custom_llm_provider,
             extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
+            optional_params=_rust_ocr_optional_params(prepared.body),
             timeout=prepared_request.effective_timeout,
         )
     except Exception as error:
@@ -349,13 +368,13 @@ async def _run_rust_aocr(
     )
     try:
         rust_response: Final = await rust_ocr_bridge.aocr(
-            model=prepared_request.model,
-            document=prepared_request.document,
+            model=_rust_ocr_model(prepared.body),
+            document=_rust_ocr_document(prepared.body),
             api_key=prepared.api_key,
             api_base=prepared.api_base,
             custom_llm_provider=prepared_request.custom_llm_provider,
             extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
+            optional_params=_rust_ocr_optional_params(prepared.body),
             timeout=prepared_request.effective_timeout,
         )
     except Exception as error:

litellm/rust_bridge/ocr.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from typing import Final, Protocol, cast  # noqa: TID251  # native extension exposes dynamically typed callables
 
 import httpx
@@ -61,17 +61,40 @@ def load_rust_aocr() -> RustAocr | None:
     return _AOCR.load()
 
 
+def _public_ocr(arguments: dict[str, object]) -> object:
+    import litellm
+
+    public_ocr: Final = cast(  # cast-ok: the decorated public OCR callable accepts the retained argument bag
+        Callable[..., object], litellm.ocr
+    )
+    return public_ocr(**arguments)
+
+
+async def _public_aocr(arguments: dict[str, object]) -> object:
+    import litellm
+
+    public_aocr: Final = cast(  # cast-ok: the decorated public OCR callable accepts the retained argument bag
+        Callable[..., Awaitable[object]], litellm.aocr
+    )
+    return await public_aocr(**arguments)
+
+
 def ocr(
+    arguments: dict[str, object] | None = None,
     *,
-    model: str,
-    document: dict[str, object],
-    api_key: str | None,
-    api_base: str | None,
-    custom_llm_provider: str | None,
-    extra_headers: dict[str, object] | None,
-    optional_params: dict[str, object],
-    timeout: float | httpx.Timeout | None,
-) -> dict[str, object] | None:
+    model: str | None = None,
+    document: dict[str, object] | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    custom_llm_provider: str | None = None,
+    extra_headers: dict[str, object] | None = None,
+    optional_params: dict[str, object] | None = None,
+    timeout: float | httpx.Timeout | None = None,
+) -> object:
+    if arguments is not None:
+        return _public_ocr(arguments)
+    if model is None or document is None or optional_params is None:
+        raise TypeError("Native OCR requires model, document, and optional_params")
     rust_ocr: Final = load_rust_ocr()
     if rust_ocr is None:
         return None
@@ -88,16 +111,21 @@ def ocr(
 
 
 async def aocr(
+    arguments: dict[str, object] | None = None,
     *,
-    model: str,
-    document: dict[str, object],
-    api_key: str | None,
-    api_base: str | None,
-    custom_llm_provider: str | None,
-    extra_headers: dict[str, object] | None,
-    optional_params: dict[str, object],
-    timeout: float | httpx.Timeout | None,
-) -> dict[str, object] | None:
+    model: str | None = None,
+    document: dict[str, object] | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    custom_llm_provider: str | None = None,
+    extra_headers: dict[str, object] | None = None,
+    optional_params: dict[str, object] | None = None,
+    timeout: float | httpx.Timeout | None = None,
+) -> object:
+    if arguments is not None:
+        return await _public_aocr(arguments)
+    if model is None or document is None or optional_params is None:
+        raise TypeError("Native OCR requires model, document, and optional_params")
     rust_aocr: Final = load_rust_aocr()
     if rust_aocr is None:
         return None

tests/test_litellm_rust/README.md

@@ -10,4 +10,4 @@ Run `make test-rust-extension` as the acceptance command. It builds a fresh whee
 
 Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs
 
-The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible
+The OCR callback contracts and supported request behavior are strict. Five request and guardrail cases remain non-strict expected failures for post-call route attribution, sanitized provider errors, timeout-specific mapping, and native Azure validation

tests/test_litellm_rust/conftest.py

@@ -30,9 +30,22 @@ CALLBACK_ATTRIBUTES: Final = (
     "_async_failure_callback",
 )
 EXPECTED_FAILURE_REASONS: Final = {
-    "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070",
-    "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070",
-    "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070",
+    (
+        "ocr/test_guardrails.py",
+        "test_native_aocr_post_call_content_filter_blocks_matching_markdown",
+    ): "requires OCR post-call guardrail route attribution",
+    (
+        "ocr/test_requests.py",
+        "test_native_ocr_maps_provider_400_without_exposing_response_body",
+    ): "requires sanitized native OCR provider errors",
+    (
+        "ocr/test_requests.py",
+        "test_native_ocr_raises_transport_error_when_request_exceeds_timeout",
+    ): "requires timeout-specific native OCR exception mapping",
+    (
+        "ocr/test_requests.py",
+        "test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks",
+    ): "requires native Azure validation before host authentication",
 }
 
 
@@ -98,7 +111,8 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
         if "test_litellm_rust" not in item.path.parts:
             continue
         relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :])
-        reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path)
+        test_name: Final = item.name.partition("[")[0]
+        reason: Final = EXPECTED_FAILURE_REASONS.get((relative_path, test_name))
         if reason is not None:
             item.add_marker(pytest.mark.xfail(reason=reason, strict=False))
 

tests/test_litellm_rust/ocr/test_callbacks.py

@@ -10,6 +10,7 @@ import litellm
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
 from tests.test_litellm_rust.support.requests import (
     OCR_DOCUMENT,
     OCR_RESPONSE,
@@ -18,7 +19,6 @@ from tests.test_litellm_rust.support.requests import (
     request_body,
     request_headers,
 )
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
 
 pytestmark = pytest.mark.requires_rust_extension
 
@@ -41,7 +41,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_
     observations: Final = []
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             observations.append((model, copy.deepcopy(kwargs["additional_args"])))
 
     call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0])
@@ -64,13 +64,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
     observed: Final = []
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             request_body(kwargs)["include_image_base64"] = True
             if raise_after_edit:
                 raise RuntimeError("pre-call callback failed")
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             observed.append(copy.deepcopy(request_body(kwargs)))
 
     call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False)
@@ -83,11 +83,11 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_
     observed: Final = []
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             request_headers(kwargs)["x-audit-tag"] = "reviewed"
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             observed.append(dict(request_headers(kwargs)))
 
     call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()])
@@ -107,12 +107,12 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_
     aliases: Final = []
 
     class Retain(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             aliases.append(request_body(kwargs)["document"] is original)
             retained.append(request_body(kwargs)["document"])
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             original["document_url"] = replacement_url
 
     arguments: Final = {
@@ -143,7 +143,7 @@ def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_docum
     retained: Final = []
 
     class RetainAndReplace(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             body = request_body(kwargs)
             retained.append(body["document"])
             body["document"] = replacement
@@ -165,11 +165,11 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov
     observed: Final = []
 
     class Rebind(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             observed.append(request_body(kwargs))
 
     call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()])
@@ -182,11 +182,11 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_
     queued: Final = []
 
     class QueuePayload(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             queued.append(request_body(kwargs))
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             request_body(kwargs)["queued-edit"] = True
 
     call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()])
@@ -200,7 +200,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
     finished: Final = threading.Event()
 
     class Stash(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             kwargs["test-token"] = token
 
         def log_success_event(self, kwargs, response_obj, start_time, end_time):
@@ -280,7 +280,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal
     observed: Final = []
 
     class TrackInFlightRequest(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             kwargs["request-token"] = token
 
         def log_failure_event(self, kwargs, response_obj, start_time, end_time):
@@ -364,7 +364,7 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context
             return "caller-token"
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, _messages, kwargs):
+        def log_pre_api_call(self, model, messages, kwargs):
             assert request_headers(kwargs)["Authorization"] == "Bearer caller-token"
             observations.append("pre_call")
             request_headers(kwargs)["Authorization"] = "Bearer edited"
@@ -465,6 +465,7 @@ async def test_native_azure_ocr_releases_token_provider_after_terminal_outcome(
 ) -> None:
     import gc
     import weakref
+
     from tests.test_litellm_rust.support.callback_recorder import drain_logging
     class Provider:
         def __call__(self) -> str:

tests/test_litellm_rust/support/callback_recorder.py

@@ -81,7 +81,7 @@ class RecordingLogger(CustomLogger):
         await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout)
         return tuple(event for event in self.events if event.name == name)
 
-    def log_pre_api_call(self, model, _messages, kwargs):
+    def log_pre_api_call(self, model, messages, kwargs):
         self._record("log_pre_api_call", kwargs)
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):