test(ocr): assert explicit backend dispatch

tests/test_litellm_rust/README.md

@@ -4,10 +4,10 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be
 
 A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions
 
-`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `test_ocr.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch
+`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport
 
 Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports
 
-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 verifies the selected transport at the wire boundary
+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

tests/test_litellm_rust/ocr/test_dispatch.py

@@ -0,0 +1,44 @@
+from typing import Final
+from unittest.mock import Mock
+
+import pytest
+
+import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from litellm.ocr import main as ocr_main
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def ocr_server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
+    return recording_server
+
+
+@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"])
+def test_public_ocr_dispatches_according_to_rust_setting(
+    ocr_server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    rust_enabled: bool,
+) -> None:
+    rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr)
+    python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr)
+    monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call)
+    monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call)
+    litellm.rust(rust_enabled)
+
+    response: Final = litellm.ocr(
+        model=OCR_MODEL,
+        document=OCR_DOCUMENT,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert isinstance(response, OCRResponse)
+    assert response.pages[0].markdown == "native OCR response"
+    assert rust_call.call_count == int(rust_enabled)
+    assert python_call.call_count == int(not rust_enabled)
+    assert len(ocr_server.requests) == 1

tests/test_litellm_rust/support/response_marker.py

@@ -1,4 +0,0 @@
-def has_rust_response_marker(response: object) -> bool:
-    from litellm.rust_bridge.provenance import has_rust_response_marker as implementation
-
-    return implementation(response)

tests/test_litellm_rust/test_ocr.py

@@ -7,14 +7,13 @@ from typing import Final
 
 import pytest
 
-import litellm
+from litellm.rust_bridge import ocr as rust_ocr_bridge
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
 @dataclass(frozen=True, slots=True)
 class RecordedOCRRequest:
-    headers: dict[str, str]
     body: object
 
 
@@ -26,7 +25,6 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest
         def do_POST(self) -> None:
             requests.append(
                 RecordedOCRRequest(
-                    headers={name.lower(): value for name, value in self.headers.items()},
                     body=json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
                 )
             )
@@ -57,26 +55,28 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest
         thread.join()
 
 
-@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"])
-def test_public_ocr_dispatches_according_to_rust_setting(
-    ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], rust_enabled: bool
+def test_native_ocr_with_compiled_rust_extension(
+    ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]],
 ) -> None:
-    litellm.rust(rust_enabled)
     server, requests = ocr_server
     address: Final = server.server_address
     host: Final = str(address[0])
     port: Final = int(address[1])
 
-    response: Final = litellm.ocr(
-        model="mistral/mistral-ocr-latest",
+    response: Final = rust_ocr_bridge.ocr(
+        model="mistral-ocr-latest",
         document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
         api_key="test-key",
         api_base=f"http://{host}:{port}",
+        custom_llm_provider="mistral",
+        extra_headers=None,
+        optional_params={},
+        timeout=None,
     )
 
-    assert response.pages[0].markdown == "native OCR response"
+    assert response is not None
+    assert response["pages"][0]["markdown"] == "native OCR response"
     assert len(requests) == 1
-    assert ("user-agent" in requests[0].headers) == (not rust_enabled)
     assert requests[0].body == {
         "model": "mistral-ocr-latest",
         "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},