fix: align rust OCR request preparation

litellm/ocr/main.py

@@ -53,7 +53,6 @@ class _PreparedOCRRequest:
 class _PreparedRustOCRCall:
     api_key: Optional[str]
     headers: dict[str, object]
-    complete_url: str
 
 
 def _timeout_to_seconds(
@@ -211,7 +210,6 @@ def _prepare_rust_ocr_call(
     return _PreparedRustOCRCall(
         api_key=resolved_api_key,
         headers=cast(dict[str, object], resolved_headers),
-        complete_url=resolved_complete_url,
     )
 
 
@@ -239,7 +237,7 @@ def _run_rust_ocr(
             api_key=prepared.api_key,
             api_base=prepared_request.api_base,
             custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared_request.extra_headers,
+            extra_headers=prepared.headers,
             optional_params=prepared_request.optional_params,
             timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
         )
@@ -262,7 +260,7 @@ async def _run_rust_aocr(
             api_key=prepared.api_key,
             api_base=prepared_request.api_base,
             custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared_request.extra_headers,
+            extra_headers=prepared.headers,
             optional_params=prepared_request.optional_params,
             timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
         )
@@ -350,6 +348,9 @@ async def aocr(
             extra_headers=extra_headers,
             kwargs=kwargs,
         )
+        model = prepared.model
+        custom_llm_provider = prepared.custom_llm_provider
+        local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
         if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled():
             rust_aocr = load_rust_aocr()
@@ -487,6 +488,9 @@ def ocr(
             extra_headers=extra_headers,
             timeout=timeout,
         )
+        model = prepared.model
+        custom_llm_provider = prepared.custom_llm_provider
+        local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
         # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge.
         if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled():

tests/test_litellm/ocr/test_rust_bridge.py

@@ -32,6 +32,10 @@ FAKE_OCR_RESPONSE: dict[str, object] = {
 }
 
 
+class CapturedException(Exception):
+    pass
+
+
 class RecordingBridge:
     """A fake ``RustOcr`` callable that records the args it was handed."""
 
@@ -96,6 +100,36 @@ class RecordingAsyncBridge:
         return dict(FAKE_OCR_RESPONSE)
 
 
+class RaisingBridge:
+    def __call__(
+        self,
+        model: str,
+        document: dict[str, object],
+        api_key: str | None,
+        api_base: str | None,
+        custom_llm_provider: str,
+        extra_headers: dict[str, object] | None,
+        optional_params: dict[str, object],
+        timeout_seconds: float | None,
+    ) -> dict[str, object]:
+        raise RuntimeError("bridge failed")
+
+
+class RaisingAsyncBridge:
+    async def __call__(
+        self,
+        model: str,
+        document: dict[str, object],
+        api_key: str | None,
+        api_base: str | None,
+        custom_llm_provider: str,
+        extra_headers: dict[str, object] | None,
+        optional_params: dict[str, object],
+        timeout_seconds: float | None,
+    ) -> dict[str, object]:
+        raise RuntimeError("bridge failed")
+
+
 class RecordingLogging:
     """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``."""
 
@@ -302,7 +336,10 @@ def test_run_rust_ocr_forwards_args_and_wraps_response():
         "api_key": "sk-test",
         "api_base": "https://proxy.internal",
         "custom_llm_provider": "mistral",
-        "extra_headers": {"x-trace-id": "trace-1"},
+        "extra_headers": {
+            "Authorization": "Bearer sk-test",
+            "x-trace-id": "trace-1",
+        },
         "optional_params": {"include_image_base64": True},
         "timeout_seconds": 12.5,
     }
@@ -413,11 +450,33 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge):
     assert call["document"] == DOCUMENT
     assert call["api_key"] == "sk-test"
     assert call["custom_llm_provider"] == "mistral"
-    assert call["extra_headers"] == {"x-trace-id": "trace-1"}
+    assert call["extra_headers"] == {
+        "Authorization": "Bearer sk-test",
+        "x-trace-id": "trace-1",
+    }
     # Raw OCR params ride along in optional_params; Rust filters to supported keys.
     assert call["optional_params"].get("include_image_base64") is True
 
 
+def test_ocr_exception_type_uses_resolved_provider_context(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    captured: dict[str, object] = {}
+
+    def fake_exception_type(**kwargs: object) -> CapturedException:
+        captured.update(kwargs)
+        return CapturedException("wrapped")
+
+    monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
+    litellm.use_litellm_rust(True, ocr=RaisingBridge())
+
+    with pytest.raises(CapturedException):
+        litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
+
+    assert captured["model"] == "mistral-ocr-latest"
+    assert captured["custom_llm_provider"] == "mistral"
+
+
 @pytest.mark.asyncio
 async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge):
     response = await litellm.aocr(
@@ -436,10 +495,33 @@ async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge):
     assert call["document"] == DOCUMENT
     assert call["api_key"] == "sk-test"
     assert call["custom_llm_provider"] == "mistral"
-    assert call["extra_headers"] == {"x-trace-id": "trace-1"}
+    assert call["extra_headers"] == {
+        "Authorization": "Bearer sk-test",
+        "x-trace-id": "trace-1",
+    }
     assert call["optional_params"].get("include_image_base64") is True
 
 
+@pytest.mark.asyncio
+async def test_aocr_exception_type_uses_resolved_provider_context(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    captured: dict[str, object] = {}
+
+    def fake_exception_type(**kwargs: object) -> CapturedException:
+        captured.update(kwargs)
+        return CapturedException("wrapped")
+
+    monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
+    litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge())
+
+    with pytest.raises(CapturedException):
+        await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
+
+    assert captured["model"] == "mistral-ocr-latest"
+    assert captured["custom_llm_provider"] == "mistral"
+
+
 def test_ocr_forwards_timeout_to_rust(fake_bridge):
     """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s
     client ceiling doesn't silently override shorter deadlines."""