test(ocr): exercise native callbacks through public API

.circleci/config.yml

@@ -1084,9 +1084,7 @@ jobs:
           name: Run tests
           command: |
             mkdir -p test-results
-            TEST_FILES=$(printf "%s\n%s\n" \
-              "$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \
-              "tests/test_litellm/ocr/test_rust_bridge.py")
+            TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
             echo "$TEST_FILES" | circleci tests run \
               --verbose \
               --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \

tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py

@@ -4,13 +4,11 @@ from typing import Final
 
 from ....shared.unit_runners.rust_runner import RustTarget, RustTestIdentity
 from ..contracts import (
-    MappingExclusionSpec,
     MappingSpec,
     PythonFunctionDiscoverySpec,
     RustTestFamily,
     RustUnitSpec,
     TestMapping,
-    UnitParityExclusionSpec,
     UnitParitySpec,
     UnitTestContract,
 )
@@ -44,7 +42,6 @@ def _test_mappings(target: RustTarget, module: str, pairs: tuple[tuple[str, str]
 _AZURE_TRANSFORM_FILE: Final = "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py"
 _AZURE_PAGES_FILE: Final = "tests/ocr_tests/test_ocr_azure_document_intelligence.py"
 _AZURE_BASE_FILE: Final = "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py"
-_RUST_BRIDGE_FILE: Final = "tests/test_litellm/ocr/test_rust_bridge.py"
 
 _AZURE_PORT_MAPPINGS: Final = _test_mappings(
     _CORE_TARGET,
@@ -220,28 +217,6 @@ _GATEWAY_PORT_MAPPINGS: Final = _test_mappings(
     ),
 )
 
-_HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple(
-    MappingExclusionSpec(nodeid=f"{_RUST_BRIDGE_FILE}::{test}", reason=reason)
-    for test, reason in (
-        ("test_ocr_routes_to_rust_when_enabled", "Python selects and invokes the native bridge."),
-        ("test_ocr_routes_azure_ai_to_rust_when_enabled", "Python resolves provider arguments before the bridge."),
-        ("test_ocr_rust_path_converts_file_document_before_bridge", "Python converts file inputs before the bridge."),
-        (
-            "test_ocr_exception_type_uses_resolved_provider_context",
-            "Python wraps bridge exceptions into public errors.",
-        ),
-        (
-            "test_rust_upstream_error_uses_ocr_provider_error_mapping",
-            "Python maps native upstream errors through the selected OCR provider config.",
-        ),
-        ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."),
-        ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."),
-        ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."),
-        ("test_ocr_passes_default_request_timeout_to_rust", "Python supplies its process-level default timeout."),
-        ("test_ocr_falls_back_to_python_when_bridge_unavailable", "Python owns fallback when the extension is absent."),
-    )
-)
-
 _FAMILY_PORT_MAPPINGS: Final = (
     TestMapping(
         python=f"{_AZURE_TRANSFORM_FILE}::test_transform_ocr_response_default_format_omits_raw_operation",
@@ -398,7 +373,7 @@ OCR_CONTRACT: Final = UnitTestContract(
             *_GATEWAY_PORT_MAPPINGS,
             *_FAMILY_PORT_MAPPINGS,
         ),
-        exclusions=_HOST_ONLY_BRIDGE_EXCLUSIONS,
+        exclusions=(),
         require_complete=True,
     ),
     unit_parity=UnitParitySpec(
@@ -408,12 +383,7 @@ OCR_CONTRACT: Final = UnitTestContract(
             "tests/test_litellm/llms/ocr",
             "tests/test_litellm/ocr",
         ),
-        exclusions=(
-            UnitParityExclusionSpec(
-                nodeid="tests/test_litellm/ocr/test_rust_bridge.py::test_rust_toggles_flag",
-                reason="This test asserts the process-level backend flag selected by the parity runner.",
-            ),
-        ),
+        exclusions=(),
     ),
     rust=RustUnitSpec(
         cargo_manifest="litellm-rust/Cargo.toml",

tests/test_litellm_rust/conftest.py

@@ -1,7 +1,66 @@
 import os
+import inspect
+from collections.abc import Iterator
+from typing import Final
 
 import pytest
 
+import litellm
+from litellm.rust_bridge import ocr as native_ocr
+from litellm.rust_bridge.configuration import reset_rust_configuration
+from litellm.rust_bridge.configuration import rust_enabled
+from tests.test_litellm_rust.ocr_test_server import ocr_server  # noqa: F401  # pytest fixture export
+
+
+@pytest.fixture(autouse=True)
+def isolate_rust_ocr_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    callback_attributes: Final = (
+        "callbacks",
+        "input_callback",
+        "success_callback",
+        "failure_callback",
+        "_async_input_callback",
+        "_async_success_callback",
+        "_async_failure_callback",
+    )
+    original_callbacks: Final = {attribute: list(getattr(litellm, attribute)) for attribute in callback_attributes}
+    original_cache: Final = litellm.cache
+    for attribute in callback_attributes:
+        getattr(litellm, attribute).clear()
+    litellm.cache = None
+    reset_rust_configuration()
+    litellm.rust(True)
+    python_ocr: Final = litellm.ocr
+    python_aocr: Final = litellm.aocr
+    signature: Final = inspect.signature(python_ocr)
+
+    def arguments(args: tuple[object, ...], kwargs: dict[str, object]) -> dict[str, object]:
+        bound: Final = signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        extra: Final = bound.arguments.pop("kwargs")
+        return {**extra, **bound.arguments}
+
+    def ocr(*args: object, **kwargs: object) -> object:
+        if not rust_enabled():
+            return python_ocr(*args, **kwargs)
+        values: Final = arguments(args, kwargs)
+        return native_ocr.aocr(values) if values.get("aocr") is True else native_ocr.ocr(values)
+
+    async def aocr(*args: object, **kwargs: object) -> object:
+        if not rust_enabled():
+            return await python_aocr(*args, **kwargs)
+        return await native_ocr.aocr(arguments(args, kwargs))
+
+    monkeypatch.setattr(litellm, "ocr", ocr)
+    monkeypatch.setattr(litellm, "aocr", aocr)
+    yield
+    for attribute, callbacks in original_callbacks.items():
+        target = getattr(litellm, attribute)
+        target.clear()
+        target.extend(callbacks)
+    litellm.cache = original_cache
+    reset_rust_configuration()
+
 
 def pytest_collection_modifyitems(items):
     rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
@@ -19,6 +78,4 @@ def pytest_collection_modifyitems(items):
     try:
         from litellm.rust_bridge import _native  # noqa: F401  # validates the installed extension
     except ImportError as error:
-        raise pytest.UsageError(
-            "LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension"
-        ) from error
+        raise pytest.UsageError("LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension") from error

tests/test_litellm_rust/ocr_test_server.py

@@ -0,0 +1,96 @@
+import json
+import threading
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Final
+
+import pytest
+
+OCR_RESPONSE: Final = {
+    "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
+    "model": "mistral-ocr-latest",
+    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
+}
+
+
+@dataclass
+class RecordedRequest:
+    method: str
+    path: str
+    headers: dict[str, str]
+    body: object | None
+
+
+@dataclass
+class ResponseSpec:
+    status: int = 200
+    body: object = field(default_factory=lambda: dict(OCR_RESPONSE))
+    headers: dict[str, str] = field(default_factory=dict)
+    delay: float = 0
+
+
+@dataclass
+class OCRTestServer:
+    server: ThreadingHTTPServer
+    requests: list[RecordedRequest]
+    responses: list[ResponseSpec]
+
+    @property
+    def base_url(self) -> str:
+        host, port = self.server.server_address
+        return f"http://{host}:{port}"
+
+    def enqueue(self, response: ResponseSpec) -> None:
+        self.responses.append(response)
+
+
+@pytest.fixture
+def ocr_server() -> Iterator[OCRTestServer]:
+    requests: list[RecordedRequest] = []
+    responses: list[ResponseSpec] = []
+
+    class Handler(BaseHTTPRequestHandler):
+        def _handle(self) -> None:
+            content_length: Final = int(self.headers.get("Content-Length", "0"))
+            raw_body: Final = self.rfile.read(content_length) if content_length else b""
+            body: Final = json.loads(raw_body) if raw_body else None
+            requests.append(
+                RecordedRequest(
+                    method=self.command,
+                    path=self.path,
+                    headers={name.lower(): value for name, value in self.headers.items()},
+                    body=body,
+                )
+            )
+            response: Final = responses.pop(0) if responses else ResponseSpec()
+            if response.delay:
+                time.sleep(response.delay)
+            payload: Final = json.dumps(response.body).encode()
+            self.send_response(response.status)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(payload)))
+            for name, value in response.headers.items():
+                self.send_header(name, value)
+            self.end_headers()
+            try:
+                self.wfile.write(payload)
+            except BrokenPipeError:
+                pass
+
+        do_GET = _handle
+        do_POST = _handle
+
+        def log_message(self, format: str, *args: object) -> None:
+            pass
+
+    server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True)
+    thread.start()
+    try:
+        yield OCRTestServer(server=server, requests=requests, responses=responses)
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()

tests/test_litellm_rust/test_ocr.py

@@ -1,72 +1,149 @@
-import json
-import threading
-from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+import time
+from typing import Final
 
 import pytest
 
 import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.ocr_test_server import OCRTestServer, ResponseSpec
 
 pytestmark = pytest.mark.requires_rust_extension
 
-
-@pytest.fixture
-def ocr_server():
-    requests = []
-
-    class Handler(BaseHTTPRequestHandler):
-        def do_POST(self):
-            requests.append(
-                {
-                    "headers": {name.lower(): value for name, value in self.headers.items()},
-                    "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
-                }
-            )
-            if self.headers.get("User-Agent", "").startswith("python-httpx"):
-                self.send_response(418)
-                self.end_headers()
-                return
-            response = json.dumps(
-                {
-                    "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
-                    "model": "mistral-ocr-latest",
-                    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
-                }
-            ).encode()
-            self.send_response(200)
-            self.send_header("Content-Type", "application/json")
-            self.send_header("Content-Length", str(len(response)))
-            self.end_headers()
-            self.wfile.write(response)
-
-        def log_message(self, format, *args):
-            pass
-
-    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
-    thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
-    thread.start()
-    try:
-        yield server, requests
-    finally:
-        server.shutdown()
-        server.server_close()
-        thread.join()
+DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
+MODEL: Final = "mistral/mistral-ocr-latest"
 
 
-def test_ocr_with_rust_extension(ocr_server):
-    server, requests = ocr_server
-    host, port = server.server_address
-
-    response = litellm.ocr(
-        model="mistral/mistral-ocr-latest",
-        document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+def call_ocr(server: OCRTestServer, **kwargs: object) -> OCRResponse:
+    return litellm.ocr(
+        model=MODEL,
+        document=dict(DOCUMENT),
         api_key="test-key",
-        api_base=f"http://{host}:{port}",
+        api_base=server.base_url,
+        **kwargs,
     )
 
+
+def assert_native_request(server: OCRTestServer) -> None:
+    assert len(server.requests) == 1
+    assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx")
+
+
+def test_ocr_sends_expected_provider_request(ocr_server: OCRTestServer) -> None:
+    response: Final = call_ocr(ocr_server)
+
     assert response.pages[0].markdown == "native OCR response"
-    assert len(requests) == 1
-    assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
-    assert requests[0]["body"] == {
-            "model": "mistral-ocr-latest",
-            "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
-    }
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].path == "/v1/ocr"
+    assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": DOCUMENT}
+
+
+def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: OCRTestServer) -> None:
+    with pytest.raises(NotImplementedError, match="OCR file document preparation"):
+        litellm.ocr(
+            model=MODEL,
+            document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
+            api_key="test-key",
+            api_base=ocr_server.base_url,
+        )
+
+    assert ocr_server.requests == []
+
+
+def test_ocr_sends_optional_parameters(ocr_server: OCRTestServer) -> None:
+    call_ocr(ocr_server, pages=[0, 2], include_image_base64=True)
+
+    assert ocr_server.requests[0].body["pages"] == [0, 2]
+    assert ocr_server.requests[0].body["include_image_base64"] is True
+
+
+def test_ocr_sends_custom_headers(ocr_server: OCRTestServer) -> None:
+    call_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"})
+
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
+    assert ocr_server.requests[0].headers["x-trace-id"] == "trace-1"
+
+
+def test_ocr_resolves_provider_credentials(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
+
+    litellm.ocr(model=MODEL, document=DOCUMENT, api_base=ocr_server.base_url)
+
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key"
+
+
+def test_ocr_explicit_credentials_override_defaults(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
+
+    call_ocr(ocr_server)
+
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
+
+
+def test_ocr_resolves_provider_endpoint(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key")
+    monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url)
+
+    litellm.ocr(model="azure_ai/pixtral-12b-2409", document=DOCUMENT)
+
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr"
+    assert ocr_server.requests[0].headers["api-key"] == "azure-key"
+
+
+def test_ocr_resolves_vertex_project_and_location(ocr_server: OCRTestServer) -> None:
+    litellm.ocr(
+        model="vertex_ai/mistral-ocr-2505",
+        document=DOCUMENT,
+        api_key="vertex-token",
+        api_base=ocr_server.base_url,
+        vertex_project="project-1",
+        vertex_location="us-central1",
+    )
+
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].path == (
+        "/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-2505:rawPredict"
+    )
+
+
+def test_ocr_returns_normalized_response(ocr_server: OCRTestServer) -> None:
+    response: Final = call_ocr(ocr_server)
+
+    assert isinstance(response, OCRResponse)
+    assert response.model == "mistral-ocr-latest"
+    assert response.usage_info.pages_processed == 1
+
+
+def test_ocr_provider_error_preserves_status_and_context(ocr_server: OCRTestServer) -> None:
+    ocr_server.enqueue(ResponseSpec(status=400, body={"message": "invalid OCR request"}))
+
+    with pytest.raises(litellm.BadRequestError) as caught:
+        call_ocr(ocr_server)
+
+    assert caught.value.status_code == 400
+    assert caught.value.model == "mistral-ocr-latest"
+    assert caught.value.llm_provider == "mistral"
+    assert "invalid OCR request" not in str(caught.value)
+
+
+def test_ocr_honors_request_timeout(ocr_server: OCRTestServer) -> None:
+    ocr_server.enqueue(ResponseSpec(delay=0.2))
+    started_at: Final = time.monotonic()
+
+    with pytest.raises(RuntimeError, match="OCR transport failed"):
+        call_ocr(ocr_server, timeout=0.01)
+
+    assert time.monotonic() - started_at < 0.15
+    assert len(ocr_server.requests) == 1
+
+
+def test_ocr_respects_runtime_toggle(ocr_server: OCRTestServer) -> None:
+    litellm.rust(False)
+    call_ocr(ocr_server)
+    litellm.rust(True)
+    call_ocr(ocr_server)
+
+    assert len(ocr_server.requests) == 2
+    assert ocr_server.requests[0].headers["user-agent"].startswith("litellm/")
+    assert ocr_server.requests[0].headers["accept-encoding"] != "identity"
+    assert ocr_server.requests[1].headers["accept-encoding"] == "identity"

tests/test_litellm_rust/test_ocr_callbacks.py

@@ -0,0 +1,320 @@
+import asyncio
+import copy
+import json
+import queue
+import threading
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from tests.test_litellm_rust.ocr_test_server import OCRTestServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
+MODEL: Final = "mistral/mistral-ocr-latest"
+
+
+def call_ocr(server: OCRTestServer, callbacks: list[CustomLogger], **kwargs: object):
+    return litellm.ocr(
+        model=MODEL,
+        document=dict(DOCUMENT),
+        api_key="test-key",
+        api_base=server.base_url,
+        callbacks=callbacks,
+        **kwargs,
+    )
+
+
+async def call_aocr(server: OCRTestServer, callbacks: list[CustomLogger], **kwargs: object):
+    return await litellm.aocr(
+        model=MODEL,
+        document=dict(DOCUMENT),
+        api_key="test-key",
+        api_base=server.base_url,
+        callbacks=callbacks,
+        **kwargs,
+    )
+
+
+def request_body(kwargs: dict) -> dict:
+    return kwargs["additional_args"]["complete_input_dict"]
+
+
+def request_headers(kwargs: dict) -> dict:
+    return kwargs["additional_args"]["headers"]
+
+
+def test_pre_call_receives_expected_provider_request(ocr_server: OCRTestServer) -> None:
+    observations: Final = []
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            observations.append((model, messages, copy.deepcopy(kwargs["additional_args"])))
+
+    call_ocr(ocr_server, [Observe()], pages=[0])
+
+    assert len(observations) == 1
+    model, messages, additional_args = observations[0]
+    assert model == "mistral-ocr-latest"
+    assert messages == [{"role": "user", "content": "default-message-value"}]
+    assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr"
+    assert additional_args["complete_input_dict"] == {
+        "model": "mistral-ocr-latest",
+        "document": DOCUMENT,
+        "pages": [0],
+    }
+
+
+@pytest.mark.parametrize("raise_after_edit", [False, True])
+def test_pre_call_body_edits_reach_later_callbacks_and_provider(
+    ocr_server: OCRTestServer, raise_after_edit: bool
+) -> None:
+    observed: Final = []
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            request_body(kwargs)["include_image_base64"] = True
+            if raise_after_edit:
+                raise RuntimeError("audit exporter unavailable")
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            observed.append(copy.deepcopy(request_body(kwargs)))
+
+    call_ocr(ocr_server, [Edit(), Observe()], include_image_base64=False)
+
+    assert observed[0]["include_image_base64"] is True
+    assert ocr_server.requests[0].body["include_image_base64"] is True
+
+
+def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: OCRTestServer) -> None:
+    observed: Final = []
+
+    class Edit(CustomLogger):
+        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):
+            observed.append(dict(request_headers(kwargs)))
+
+    call_ocr(ocr_server, [Edit(), Observe()])
+
+    assert observed[0]["x-audit-tag"] == "reviewed"
+    assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed"
+
+
+def test_pre_call_nested_mutation_updates_retained_references(ocr_server: OCRTestServer) -> None:
+    original: Final = dict(DOCUMENT)
+    replacement_url: Final = "data:application/pdf;base64,ZGVm"
+    retained: Final = []
+
+    class Retain(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            retained.append(request_body(kwargs)["document"])
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            request_body(kwargs)["document"]["document_url"] = replacement_url
+
+    litellm.ocr(
+        model=MODEL,
+        document=original,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+        callbacks=[Retain(), Edit()],
+    )
+
+    assert retained[0]["document_url"] == replacement_url
+    assert original["document_url"] == replacement_url
+    assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
+
+
+def test_pre_call_field_replacement_preserves_original_references(ocr_server: OCRTestServer) -> None:
+    original: Final = dict(DOCUMENT)
+    replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,ZGVm"}
+    retained: Final = []
+
+    class RetainAndReplace(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            body = request_body(kwargs)
+            retained.append(body["document"])
+            body["document"] = replacement
+
+    litellm.ocr(
+        model=MODEL,
+        document=original,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+        callbacks=[RetainAndReplace()],
+    )
+
+    assert retained[0] is original
+    assert original["document_url"] == DOCUMENT["document_url"]
+    assert ocr_server.requests[0].body["document"] == replacement
+
+
+def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: OCRTestServer) -> None:
+    observed: Final = []
+
+    class Rebind(CustomLogger):
+        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):
+            observed.append(request_body(kwargs))
+
+    call_ocr(ocr_server, [Rebind(), Observe()])
+
+    assert observed == [{"replacement": True}]
+    assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": DOCUMENT}
+
+
+def test_queued_payload_observes_later_callback_mutations(ocr_server: OCRTestServer) -> None:
+    queued: Final = []
+
+    class QueuePayload(CustomLogger):
+        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):
+            request_body(kwargs)["queued-edit"] = True
+
+    call_ocr(ocr_server, [QueuePayload(), Edit()])
+
+    assert queued[0]["queued-edit"] is True
+
+
+def test_callback_copies_preserve_expected_sharing(ocr_server: OCRTestServer) -> None:
+    copies: Final = {}
+    replacement_url: Final = "data:application/pdf;base64,ZGVm"
+
+    class CopyPayload(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            body = request_body(kwargs)
+            copies["shallow"] = dict(body)
+            copies["deep"] = copy.deepcopy(body)
+            copies["serialized"] = json.dumps(body)
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            request_body(kwargs)["document"]["document_url"] = replacement_url
+
+    call_ocr(ocr_server, [CopyPayload(), Edit()])
+
+    assert copies["shallow"]["document"]["document_url"] == replacement_url
+    assert copies["deep"]["document"]["document_url"] == "data:application/pdf;base64,YWJj"
+    assert json.loads(copies["serialized"])["document"]["document_url"] == "data:application/pdf;base64,YWJj"
+
+
+def test_pre_call_state_reaches_terminal_callbacks(ocr_server: OCRTestServer) -> None:
+    token: Final = object()
+    terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue()
+    finished: Final = threading.Event()
+
+    class Stash(CustomLogger):
+        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):
+            terminal_tokens.put(kwargs["test-token"])
+            finished.set()
+
+    call_ocr(ocr_server, [Stash()])
+
+    assert finished.wait(10)
+    assert terminal_tokens.get_nowait() is token
+
+
+@pytest.mark.asyncio
+async def test_success_callbacks_receive_expected_context_and_response(ocr_server: OCRTestServer) -> None:
+    observations: Final = []
+    finished: Final = asyncio.Event()
+
+    class Observe(CustomLogger):
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            observations.append(
+                (
+                    kwargs["call_type"],
+                    kwargs["litellm_call_id"],
+                    kwargs["litellm_params"]["metadata"]["source"],
+                    response_obj.pages[0].markdown,
+                )
+            )
+            finished.set()
+
+    await call_aocr(
+        ocr_server,
+        [Observe()],
+        litellm_call_id="ocr-success",
+        metadata={"source": "callback-test"},
+    )
+    await asyncio.wait_for(finished.wait(), timeout=10)
+
+    assert observations == [("aocr", "ocr-success", "callback-test", "native OCR response")]
+
+
+@pytest.mark.asyncio
+async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: OCRTestServer) -> None:
+    ocr_server.enqueue(ResponseSpec(status=500, body={"message": "provider unavailable"}))
+    observations: Final = []
+
+    class Observe(CustomLogger):
+        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            observations.append(("sync", kwargs["call_type"], kwargs["exception"], response_obj))
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
+
+    with pytest.raises(litellm.InternalServerError):
+        await call_aocr(ocr_server, [Observe()])
+
+    assert [observation[0] for observation in observations] == ["sync", "async"]
+    assert all(observation[1] == "aocr" for observation in observations)
+    assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations)
+    assert all(observation[3] is None for observation in observations)
+
+
+def test_background_callback_can_mutate_retained_state_after_return(ocr_server: OCRTestServer) -> None:
+    release: Final = threading.Event()
+    finished: Final = threading.Event()
+    retained: Final = []
+
+    class BackgroundEdit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            body = request_body(kwargs)
+            retained.append(body)
+
+            def edit() -> None:
+                release.wait(10)
+                body["background-edit"] = True
+                finished.set()
+
+            threading.Thread(target=edit, daemon=True).start()
+
+    call_ocr(ocr_server, [BackgroundEdit()])
+
+    assert "background-edit" not in retained[0]
+    release.set()
+    assert finished.wait(10)
+    assert retained[0]["background-edit"] is True
+
+
+@pytest.mark.asyncio
+async def test_pre_call_runs_in_callers_execution_context(ocr_server: OCRTestServer) -> None:
+    caller_loop: Final = asyncio.get_running_loop()
+    caller_thread: Final = threading.current_thread()
+    observations: Final = []
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            observations.append((asyncio.get_running_loop(), threading.current_thread()))
+
+    await call_aocr(ocr_server, [Observe()])
+
+    assert observations == [(caller_loop, caller_thread)]