test(rust): add retained callback suite as expected failures

.github/workflows/test-rust.yml

@@ -4,6 +4,7 @@ on:
   push:
     paths:
       - "litellm-rust/**"
+      - "tests/test_litellm_rust/**"
       - ".cargo/**"
       - "pyproject.toml"
       - "rust-toolchain.toml"
@@ -20,6 +21,7 @@ on:
       - "litellm_**"
     paths:
       - "litellm-rust/**"
+      - "tests/test_litellm_rust/**"
       - ".cargo/**"
       - "pyproject.toml"
       - "rust-toolchain.toml"

tests/test_litellm_rust/README.md

@@ -0,0 +1,11 @@
+# Rust bridge tests
+
+Put API-specific tests under `ocr/`, `messages/`, or `chat/`. Keep request construction, recording servers, callback recorders, and state isolation in `support/`. Integration tests and their fixtures belong in `integrations/`.
+
+Use `backend` with `"python"` and `"rust"` for parity tests. Keep callback mutation and lifecycle assertions next to the API surface that owns them. Shared route parametrization belongs in `integrations/routes.py` only when its observable contract is the same for every route.
+
+Run the suite with `LITELLM_RUST=1 uv run pytest tests/test_litellm_rust`.
+
+The newly ported suite is marked as a non-strict expected failure until the retained callback implementation from #40070 lands. Passing cases appear as XPASS so staging coverage remains visible
+
+`isolated_backend` restores the previous backend override and callback state on exit, including after exceptions or cancellation. Scopes can nest in the same task. Run backend comparisons sequentially: overlapping scopes in different tasks raise before changing process-global state. Use separate worker processes for parallel backend comparisons

tests/test_litellm_rust/chat/__init__.py

@@ -0,0 +1 @@
+

tests/test_litellm_rust/chat/test_callback_mutation.py

@@ -0,0 +1,404 @@
+import hashlib
+import hmac
+import json
+from typing import Final
+from urllib.parse import urlsplit
+
+import pytest
+
+import litellm
+from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+def _verify_sigv4(request: RecordedRequest, secret_key: str) -> None:
+    authorization: Final = request.headers["authorization"]
+    algorithm, attributes_text = authorization.split(" ", 1)
+    attributes: Final = dict(item.split("=", 1) for item in attributes_text.split(", "))
+    credential_scope: Final = attributes["Credential"].split("/", 1)[1]
+    signed_names: Final = attributes["SignedHeaders"].split(";")
+    canonical_headers: Final = "".join(f"{name}:{' '.join(request.headers[name].split())}\n" for name in signed_names)
+    parsed_path: Final = urlsplit(request.path)
+    canonical_request: Final = "\n".join(
+        (
+            request.method,
+            parsed_path.path,
+            parsed_path.query,
+            canonical_headers,
+            attributes["SignedHeaders"],
+            hashlib.sha256(request.raw_body).hexdigest(),
+        )
+    )
+    amz_date: Final = request.headers["x-amz-date"]
+    string_to_sign: Final = "\n".join(
+        (algorithm, amz_date, credential_scope, hashlib.sha256(canonical_request.encode()).hexdigest())
+    )
+    date, region, service, terminator = credential_scope.split("/")
+    date_key: Final = hmac.new(f"AWS4{secret_key}".encode(), date.encode(), hashlib.sha256).digest()
+    region_key: Final = hmac.new(date_key, region.encode(), hashlib.sha256).digest()
+    service_key: Final = hmac.new(region_key, service.encode(), hashlib.sha256).digest()
+    signing_key: Final = hmac.new(service_key, terminator.encode(), hashlib.sha256).digest()
+    expected: Final = hmac.new(signing_key, string_to_sign.encode(), hashlib.sha256).hexdigest()
+    assert hmac.compare_digest(attributes["Signature"], expected)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("provider", ["anthropic", "bedrock"])
+@pytest.mark.parametrize("rebind_logging_view", [False, True])
+@pytest.mark.parametrize("status", [200, 429])
+@pytest.mark.parametrize("native", [False, True])
+async def test_chat_retains_callback_edits_through_public_dispatch(
+    recording_server: RecordingServer,
+    asynchronous: bool,
+    provider: str,
+    rebind_logging_view: bool,
+    status: int,
+    native: bool,
+) -> None:
+    import threading
+
+    litellm.rust(native)
+
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+    from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
+
+    recording_server.default_response = ResponseSpec(
+        status=status,
+        body=(
+            MESSAGES_RESPONSE
+            if provider == "anthropic"
+            else {
+                "output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}},
+                "stopReason": "end_turn",
+                "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
+                "metrics": {"latencyMs": 1},
+            }
+        ),
+    )
+    caller_thread: Final = threading.current_thread()
+    observations: Final = []
+
+    class EditingLogger(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            super().log_pre_api_call(model, messages, kwargs)
+            body = kwargs["additional_args"]["complete_input_dict"]
+            headers = kwargs["additional_args"]["headers"]
+            if provider == "anthropic":
+                assert isinstance(body, dict)
+                body["messages"][0]["content"][0]["text"] = "edited by callback"
+                body["max_tokens"] = 32
+            else:
+                assert isinstance(body, str)
+                assert json.loads(body)["messages"][0]["content"][0]["text"] == "original"
+            headers["x-retained-callback"] = "original"
+            observations.append((threading.current_thread(), body, headers))
+            if rebind_logging_view:
+                kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
+                kwargs["additional_args"]["headers"] = {"x-retained-callback": "replacement"}
+
+    recorder: Final = EditingLogger()
+    kwargs: Final = {
+        "model": "anthropic/claude-opus-5" if provider == "anthropic" else "bedrock/anthropic.claude-opus-5",
+        "messages": [{"role": "user", "content": "original"}],
+        "max_tokens": 64,
+        "api_key": "test-key",
+        "api_base": recording_server.base_url,
+        "callbacks": [recorder],
+        "num_retries": 0,
+        **(
+            {"aws_access_key_id": "test", "aws_secret_access_key": "test", "aws_region_name": "us-east-1"}
+            if provider == "bedrock"
+            else {}
+        ),
+    }
+    if status != 200:
+        with pytest.raises((litellm.APIError, litellm.RateLimitError)) as raised:
+            await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
+        failure_event: Final = "async_log_failure_event" if asynchronous else "log_failure_event"
+        failures: Final = await recorder.wait_for_async(failure_event)
+        assert raised.value.status_code == status
+        assert failures[0].kwargs["exception"] is raised.value
+        assert recorder.names.count(failure_event) == 1
+        assert recorder.names.count("log_pre_api_call") == 1
+        assert len(recording_server.requests) == 1
+        return
+
+    response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
+    event_name: Final = "async_log_success_event" if asynchronous else "log_success_event"
+    events: Final = await recorder.wait_for_async(event_name)
+
+    if native:
+        assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert recorder.names.count(event_name) == 1
+    if asynchronous and provider == "anthropic" and not native:
+        assert observations[0][0] is not caller_thread
+    else:
+        assert observations[0][0] is caller_thread
+    assert events[0].response is response
+    assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == (
+        "edited by callback" if provider == "anthropic" else "original"
+    )
+    assert recording_server.requests[0].headers["x-retained-callback"] == "original"
+    body: Final = recording_server.requests[0].body
+    assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == (
+        32 if provider == "anthropic" else 64
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("status", [200, 429])
+@pytest.mark.parametrize("native", [False, True])
+async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
+    recording_server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    asynchronous: bool,
+    status: int,
+    native: bool,
+) -> None:
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+
+    litellm.rust(native)
+    monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+    recording_server.default_response = ResponseSpec(
+        status=status,
+        body={
+            "output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}},
+            "stopReason": "end_turn",
+            "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
+            "metrics": {"latencyMs": 1},
+        },
+    )
+    observations: Final[list[tuple[int, int]]] = []
+    replacement_body: Final = {"replacement": True}
+    rebound_headers: Final = {"x-callback-header": "rebound"}
+
+    class FirstCallback(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            additional: Final = kwargs["additional_args"]
+            callback_body: Final = additional["complete_input_dict"]
+            original_headers: Final = additional["headers"]
+            observations.append((id(kwargs), id(additional)))
+            assert isinstance(callback_body, str)
+            assert json.loads(callback_body)["messages"][0]["content"][0]["text"] == "original"
+            kwargs["callback_marker"] = "visible"
+            original_headers["x-callback-header"] = "in-place"
+            additional["complete_input_dict"] = replacement_body
+            additional["headers"] = rebound_headers
+            raise RuntimeError("first callback failure is non-blocking")
+
+    class SecondCallback(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            additional: Final = kwargs["additional_args"]
+            observations.append((id(kwargs), id(additional)))
+            assert kwargs["callback_marker"] == "visible"
+            assert additional["complete_input_dict"] is replacement_body
+            assert additional["headers"] is rebound_headers
+
+    request: Final = {
+        "model": "bedrock/anthropic.claude-opus-5",
+        "messages": [{"role": "user", "content": "original"}],
+        "max_tokens": 64,
+        "api_base": recording_server.base_url,
+        "callbacks": [FirstCallback(), SecondCallback()],
+        "num_retries": 0,
+        "aws_access_key_id": "test-access",
+        "aws_secret_access_key": "test-secret",
+        "aws_region_name": "us-east-1",
+    }
+    if status == 200:
+        await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
+    else:
+        with pytest.raises((litellm.APIError, litellm.RateLimitError)):
+            await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
+
+    assert observations[0] == observations[1]
+    recorded: Final = recording_server.requests[0]
+    assert recorded.body["messages"][0]["content"][0]["text"] == "original"
+    assert "replacement" not in recorded.body
+    assert recorded.headers["x-callback-header"] == "in-place"
+    assert recorded.headers["authorization"].startswith("AWS4-HMAC-SHA256 ")
+    _verify_sigv4(recorded, "test-secret")
+
+
+@pytest.mark.asyncio
+async def test_native_chat_pre_call_preserves_the_caller_task_and_context(
+    recording_server: RecordingServer,
+) -> None:
+    import asyncio
+    from contextvars import ContextVar
+
+    from litellm.integrations.custom_logger import CustomLogger
+    from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
+
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    marker: Final[ContextVar[str]] = ContextVar("native_chat_marker", default="missing")
+    marker.set("caller")
+    caller_task: Final = asyncio.current_task()
+    observations: Final = []
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            observations.append((asyncio.current_task(), marker.get()))
+            marker.set("callback")
+
+    response: Final = await litellm.acompletion(
+        model="anthropic/claude-opus-5",
+        messages=[{"role": "user", "content": "task context"}],
+        max_tokens=16,
+        api_key="test-key",
+        api_base=recording_server.base_url,
+        callbacks=[Observe()],
+        num_retries=0,
+    )
+
+    assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
+    assert observations == [(caller_task, "caller")]
+    assert marker.get() == "callback"
+
+
+@pytest.mark.asyncio
+async def test_native_chat_concurrent_calls_keep_callback_roots_isolated(
+    recording_server: RecordingServer,
+) -> None:
+    import asyncio
+    import threading
+
+    from litellm.integrations.custom_logger import CustomLogger
+    from tests.test_litellm_rust.support.callback_recorder import drain_logging
+    from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
+
+    call_ids: Final = tuple(f"chat-{index}" for index in range(4))
+    recording_server.expected_requests = len(call_ids)
+    accepted: Final = tuple(threading.Event() for _ in call_ids)
+    release: Final = threading.Event()
+    for signal in accepted:
+        recording_server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE, accepted=signal, release_before_response=release))
+    tokens: Final = {call_id: object() for call_id in call_ids}
+    terminal_state: Final = []
+
+    class Correlate(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            body: Final = kwargs["additional_args"]["complete_input_dict"]
+            body["messages"][0]["content"][0]["text"] = f"callback-{kwargs['litellm_call_id']}"
+            kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]]
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"]))
+
+    logger: Final = Correlate()
+
+    async def invoke(call_id: str) -> None:
+        await litellm.acompletion(
+            model="anthropic/claude-opus-5",
+            messages=[{"role": "user", "content": f"original-{call_id}"}],
+            max_tokens=16,
+            api_key="test-key",
+            api_base=recording_server.base_url,
+            callbacks=[logger],
+            litellm_call_id=call_id,
+            num_retries=0,
+        )
+
+    tasks: Final = tuple(asyncio.create_task(invoke(call_id)) for call_id in call_ids)
+    try:
+        assert all(await asyncio.gather(*(asyncio.to_thread(signal.wait, 3) for signal in accepted)))
+    finally:
+        release.set()
+    await asyncio.gather(*tasks)
+    await drain_logging()
+
+    sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests}
+    assert sent == {f"callback-{call_id}" for call_id in call_ids}
+    assert len(terminal_state) == len(call_ids)
+    assert all(token is tokens[call_id] for call_id, token in terminal_state)
+
+
+@pytest.mark.asyncio
+async def test_native_chat_callback_can_make_a_nested_native_call(
+    recording_server: RecordingServer,
+) -> None:
+    import threading
+
+    from litellm.integrations.custom_logger import CustomLogger
+    from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
+
+    recording_server.expected_requests = 2
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    nested_responses: Final = []
+    nested_started: Final = threading.Event()
+
+    class NestedCall(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            if nested_started.is_set():
+                return
+            nested_started.set()
+            nested_responses.append(
+                litellm.completion(
+                    model="anthropic/claude-opus-5",
+                    messages=[{"role": "user", "content": "nested"}],
+                    max_tokens=16,
+                    api_key="test-key",
+                    api_base=recording_server.base_url,
+                    num_retries=0,
+                )
+            )
+
+    outer: Final = await litellm.acompletion(
+        model="anthropic/claude-opus-5",
+        messages=[{"role": "user", "content": "outer"}],
+        max_tokens=16,
+        api_key="test-key",
+        api_base=recording_server.base_url,
+        callbacks=[NestedCall()],
+        num_retries=0,
+    )
+
+    assert outer._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
+    assert nested_responses[0]._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
+    sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests}
+    assert sent == {"outer", "nested"}
+
+
+@pytest.mark.asyncio
+async def test_native_chat_cancellation_during_io_does_not_publish_a_terminal(
+    recording_server: RecordingServer,
+) -> None:
+    import asyncio
+
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+    from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
+
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, delay=0.5)
+    recorder: Final = RecordingLogger()
+    task: Final = asyncio.create_task(
+        litellm.acompletion(
+            model="anthropic/claude-opus-5",
+            messages=[{"role": "user", "content": "cancel"}],
+            max_tokens=16,
+            api_key="test-key",
+            api_base=recording_server.base_url,
+            callbacks=[recorder],
+            num_retries=0,
+        )
+    )
+    async with asyncio.timeout(10):
+        while not recording_server.requests:
+            await asyncio.sleep(0.01)
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert not {
+        "log_success_event",
+        "async_log_success_event",
+        "log_failure_event",
+        "async_log_failure_event",
+    }.intersection(recorder.names)

tests/test_litellm_rust/conftest.py

@@ -1,24 +1,181 @@
+import asyncio
 import os
+from collections.abc import AsyncIterator, Generator
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack, asynccontextmanager, contextmanager
+from threading import Lock
+from types import ModuleType
+from typing import Final, Literal, cast
 
 import pytest
+import pytest_asyncio
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+import litellm
+from litellm import utils
+from litellm.litellm_core_utils import litellm_logging
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.rust_bridge.configuration import (  # pyright: ignore[reportPrivateUsage]  # preserve raw configuration state in test isolation
+    _CONFIGURATION,
+    _parse_env_bool,
+)
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry
+from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
+
+CALLBACK_ATTRIBUTES: Final = (
+    "callbacks",
+    "input_callback",
+    "success_callback",
+    "failure_callback",
+    "_async_input_callback",
+    "_async_success_callback",
+    "_async_failure_callback",
+)
+Backend = Literal["python", "rust"]
 
 
-def pytest_collection_modifyitems(items):
-    rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
-        "1",
-        "true",
-        "yes",
-        "on",
-    }
-    if not rust_enabled:
+class _BackendScopes:
+    def __init__(self) -> None:
+        self._owner: asyncio.Task[object] | None = None
+        self._lock: Final = Lock()
+
+    @contextmanager
+    def enter(self) -> Generator[None]:
+        task: Final = asyncio.current_task()
+        with self._lock:
+            previous: Final = self._owner
+            if previous is not None and previous is not task:
+                raise RuntimeError("isolated_backend scopes cannot overlap across tasks; run them sequentially")
+            self._owner = task
+        try:
+            yield
+        finally:
+            with self._lock:
+                self._owner = previous
+
+
+_BACKEND_SCOPES: Final = _BackendScopes()
+
+
+def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
+    value: Final = getattr(container, attribute)
+    if not isinstance(value, list):
+        raise AssertionError(f"{container.__name__}.{attribute} is not a list")
+    return cast(list[object], value)
+
+
+@contextmanager
+def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]:
+    source: Final = _list_attribute(container, attribute)
+    original: Final = list(source)
+    source.clear()  # mutable-ok: test isolation mutates global registries by design
+    try:
+        yield
+    finally:
+        source.clear()
+        source.extend(original)
+        setattr(container, attribute, source)
+
+
+@contextmanager
+def _rebound(container: object, attribute: str, value: object) -> Generator[None]:
+    original: Final[object] = getattr(container, attribute)
+    setattr(container, attribute, value)
+    try:
+        yield
+    finally:
+        setattr(container, attribute, original)
+
+
+@contextmanager
+def _rust_mode(enabled: bool) -> Generator[None]:
+    with _rebound(_CONFIGURATION, "override", enabled):
+        yield
+
+
+@asynccontextmanager
+async def isolated_backend(backend: Backend) -> AsyncIterator[ExitStack]:
+    with _BACKEND_SCOPES.enter():
+        async with _isolated_backend(backend) as stack:
+            yield stack
+
+
+@asynccontextmanager
+async def _isolated_backend(backend: Backend) -> AsyncIterator[ExitStack]:
+    with ExitStack() as stack:
+        for attribute in CALLBACK_ATTRIBUTES:
+            stack.enter_context(_isolated_list(litellm, attribute))
+        stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers"))  # pyright: ignore[reportPrivateUsage]  # string-name callback cache has no public accessor
+        stack.enter_context(
+            _rebound(utils, "callback_list", [])
+        )  # rebind-ok: legacy global registry mutated by set_callbacks
+        stack.enter_context(
+            _rebound(litellm, "cache", None)
+        )  # test-quality-ok: isolate the process-global cache from native extension tests
+        stack.enter_context(isolated_prometheus_registry())
+        stack.enter_context(_rust_mode(backend == "rust"))
+        executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging")
+        stack.enter_context(_rebound(utils, "executor", executor))
+        try:
+            yield stack
+        finally:
+            try:
+                await drain_logging()
+            finally:
+                await asyncio.to_thread(executor.shutdown, wait=True)
+                await GLOBAL_LOGGING_WORKER.stop()
+
+
+@pytest_asyncio.fixture(autouse=True, loop_scope="function")
+async def isolate_rust_state() -> AsyncIterator[ExitStack]:
+    # pytest-asyncio runs fixture setup and the test body in different tasks.
+    async with _isolated_backend("rust") as stack:
+        yield stack
+
+
+@pytest.fixture
+def recording_server() -> Generator[RecordingServer]:
+    with recording_service() as server:
+        yield server
+
+
+def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
+    expected_failure = pytest.mark.xfail(
+        reason="requires the retained callback implementation from #40070",
+        strict=False,
+    )
+    for item in items:
+        if "test_litellm_rust" in item.path.parts and item.path.name != "test_ocr.py":
+            item.add_marker(expected_failure)
+
+    if not _parse_env_bool(os.environ.get("LITELLM_RUST")):
         skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
         for item in items:
-            item.add_marker(skip)
+            if "test_litellm_rust" in item.path.parts:
+                item.add_marker(skip)
         return
 
     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
+
+
+@pytest.fixture
+def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None:
+    for name in (
+        "AZURE_AI_API_KEY",
+        "AZURE_AI_API_BASE",
+        "AZURE_AD_TOKEN",
+        "AZURE_TENANT_ID",
+        "AZURE_CLIENT_ID",
+        "AZURE_CLIENT_SECRET",
+        "AZURE_USERNAME",
+        "AZURE_PASSWORD",
+    ):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setattr(litellm, "api_key", None)
+    monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False)

tests/test_litellm_rust/integrations/__init__.py

@@ -0,0 +1,75 @@
+from tests.test_litellm_rust.integrations.catalogue import (
+    DISCOVERED_ONLY_GUARDRAIL_NAMES,
+    ENTERPRISE_LOGGER_NAMES,
+    GUARDRAIL_NAMES,
+    GUARDRAIL_OBLIGATIONS,
+    LOGGER_OBLIGATIONS,
+    OSS_LOGGER_NAMES,
+    REQUIRED_GUARDRAIL_BEHAVIOR,
+    REQUIRED_LOGGER_BEHAVIOR,
+)
+from tests.test_litellm_rust.integrations.harness import (
+    AsyncBoundaryLogger,
+    GCSLiteralAIHarness,
+    GenericAPIExportHarness,
+    MutatingFailingLogger,
+    OtelHarness,
+    PurviewHarness,
+    RecordingAsyncClient,
+    RecordingVertexInstance,
+    ReviewGuardrail,
+    RunObservation,
+    gcs_literalai_harness,
+    metric_value,
+    purview_harness,
+    wait_for_audits,
+)
+from tests.test_litellm_rust.support.routes import (
+    ALL_ROUTES,
+    ASYNC_ROUTES,
+    MESSAGES_ROUTE,
+    MESSAGES_STREAM,
+    NON_STREAM_ASYNC_ROUTES,
+    OCR_ASYNC,
+    OCR_SYNC,
+    Route,
+    provider_response,
+    route_id,
+    wait_for_callback,
+)
+
+__all__ = (
+    "ALL_ROUTES",
+    "ASYNC_ROUTES",
+    "DISCOVERED_ONLY_GUARDRAIL_NAMES",
+    "ENTERPRISE_LOGGER_NAMES",
+    "GUARDRAIL_NAMES",
+    "GUARDRAIL_OBLIGATIONS",
+    "LOGGER_OBLIGATIONS",
+    "MESSAGES_ROUTE",
+    "MESSAGES_STREAM",
+    "NON_STREAM_ASYNC_ROUTES",
+    "OCR_ASYNC",
+    "OCR_SYNC",
+    "OSS_LOGGER_NAMES",
+    "REQUIRED_GUARDRAIL_BEHAVIOR",
+    "REQUIRED_LOGGER_BEHAVIOR",
+    "AsyncBoundaryLogger",
+    "GCSLiteralAIHarness",
+    "GenericAPIExportHarness",
+    "MutatingFailingLogger",
+    "OtelHarness",
+    "PurviewHarness",
+    "RecordingAsyncClient",
+    "RecordingVertexInstance",
+    "ReviewGuardrail",
+    "Route",
+    "RunObservation",
+    "gcs_literalai_harness",
+    "metric_value",
+    "provider_response",
+    "purview_harness",
+    "route_id",
+    "wait_for_audits",
+    "wait_for_callback",
+)

tests/test_litellm_rust/integrations/catalogue.py

@@ -0,0 +1,193 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final, Literal
+
+DependencyProfile = Literal["required", "enterprise"]
+
+OSS_LOGGER_NAMES: Final = frozenset(
+    {
+        "agentops",
+        "anthropic_cache_control_hook",
+        "argilla",
+        "arize",
+        "arize_phoenix",
+        "aws_sqs",
+        "azure_sentinel",
+        "azure_storage",
+        "bitbucket",
+        "braintrust",
+        "cloudzero",
+        "datadog",
+        "datadog_llm_observability",
+        "datadog_metrics",
+        "deepeval",
+        "dotprompt",
+        "dynamic_rate_limiter",
+        "dynamic_rate_limiter_v3",
+        "focus",
+        "galileo",
+        "gcs_bucket",
+        "gcs_pubsub",
+        "gitlab",
+        "humanloop",
+        "lago",
+        "langfuse",
+        "langfuse_otel",
+        "langsmith",
+        "langtrace",
+        "levo",
+        "literalai",
+        "litellm_agent",
+        "logfire",
+        "mavvrik",
+        "mlflow",
+        "newrelic",
+        "opentelemetry",
+        "openmeter",
+        "opik",
+        "otel",
+        "posthog",
+        "prometheus",
+        "s3_v2",
+        "vantage",
+        "vector_store_pre_call_hook",
+        "weave_otel",
+    }
+)
+ENTERPRISE_LOGGER_NAMES: Final = frozenset({"generic_api", "pagerduty", "resend_email", "sendgrid_email", "smtp_email"})
+GUARDRAIL_NAMES: Final = frozenset(
+    {
+        "aim",
+        "akto",
+        "alice",
+        "aporia",
+        "azure/prompt_shield",
+        "azure/text_moderations",
+        "bedrock",
+        "block_code_execution",
+        "cato_networks",
+        "cisco_ai_defense",
+        "compresr",
+        "crowdstrike_aidr",
+        "custom_code",
+        "deepkeep",
+        "dynamoai",
+        "enkryptai",
+        "generic_guardrail_api",
+        "grayswan",
+        "guardrails_ai",
+        "headroom",
+        "hiddenlayer",
+        "hide-secrets",
+        "ibm_guardrails",
+        "javelin",
+        "lakera",
+        "lakera_v2",
+        "lasso",
+        "litellm_content_filter",
+        "llm_as_a_judge",
+        "mcp_end_user_permission",
+        "mcp_jwt_signer",
+        "mcp_security",
+        "microsoft_purview",
+        "model_armor",
+        "noma",
+        "noma_v2",
+        "onyx",
+        "openai_moderation",
+        "ovalix",
+        "pangea",
+        "panw_prisma_airs",
+        "pillar",
+        "presidio",
+        "prompt_security",
+        "promptguard",
+        "qostodian_nexus",
+        "qualifire",
+        "repelloai",
+        "rubrik",
+        "semantic_guard",
+        "singulr",
+        "straiker",
+        "tool_permission",
+        "vigil_guard",
+        "xecguard",
+        "zscaler_ai_guard",
+    }
+)
+DISCOVERED_ONLY_GUARDRAIL_NAMES: Final = frozenset({"tool_policy"})
+
+
+@dataclass(frozen=True, slots=True)
+class CoverageObligation:
+    stable_name: str
+    registration_names: tuple[str, ...]
+    dependency_profile: DependencyProfile
+    behavioral_cases: tuple[str, ...]
+
+
+REQUIRED_LOGGER_BEHAVIOR: Final = frozenset({"generic_api", "gcs_bucket", "literalai", "prometheus", "opentelemetry"})
+REQUIRED_GUARDRAIL_BEHAVIOR: Final = frozenset(
+    {
+        "azure/text_moderations",
+        "crowdstrike_aidr",
+        "litellm_content_filter",
+        "microsoft_purview",
+        "rubrik",
+    }
+)
+
+LOGGER_BEHAVIORAL_CASES: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
+    {
+        "generic_api": ("generic-api-success",),
+        "gcs_bucket": ("gcs-literalai-scheduling",),
+        "literalai": ("gcs-literalai-scheduling",),
+        "prometheus": ("prometheus-string-registration",),
+        "opentelemetry": ("otel-export",),
+    }
+)
+GUARDRAIL_BEHAVIORAL_CASES: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
+    {
+        "azure/text_moderations": ("azure-text-moderation",),
+        "crowdstrike_aidr": ("crowdstrike-redaction-native-chat",),
+        "rubrik": ("rubrik-block-native-chat",),
+        "microsoft_purview": ("purview-audit-native-chat",),
+        "litellm_content_filter": ("content-filter-block",),
+    }
+)
+
+
+def _logger_obligations() -> Mapping[str, CoverageObligation]:
+    from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
+
+    registry: Final = CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE
+    expected_names: Final = OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
+    grouped: Final[dict[type[object], tuple[str, ...]]] = {
+        implementation: tuple(sorted(name for name in expected_names if registry.get(name) is implementation))
+        for implementation in frozenset(registry.values())
+    }
+    obligations: Final = {
+        names[0]: CoverageObligation(
+            stable_name=names[0],
+            registration_names=names,
+            dependency_profile="enterprise" if set(names) & ENTERPRISE_LOGGER_NAMES else "required",
+            behavioral_cases=next(
+                (LOGGER_BEHAVIORAL_CASES[name] for name in names if name in LOGGER_BEHAVIORAL_CASES), ()
+            ),
+        )
+        for names in grouped.values()
+        if names
+    }
+    missing_optional: Final = ENTERPRISE_LOGGER_NAMES - registry.keys()
+    unavailable: Final = {name: CoverageObligation(name, (name,), "enterprise", ()) for name in missing_optional}
+    return MappingProxyType({**obligations, **unavailable})
+
+
+LOGGER_OBLIGATIONS: Final = _logger_obligations()
+GUARDRAIL_OBLIGATIONS: Final = MappingProxyType(
+    {
+        name: CoverageObligation(name, (name,), "required", GUARDRAIL_BEHAVIORAL_CASES.get(name, ()))
+        for name in GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES
+    }
+)

tests/test_litellm_rust/integrations/conftest.py

@@ -0,0 +1,43 @@
+from contextlib import ExitStack
+from typing import Final
+
+import pytest
+import pytest_asyncio
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
+from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from litellm.integrations.prometheus import PrometheusLogger
+from tests.test_litellm_rust.integrations import GenericAPIExportHarness, OtelHarness, Route, provider_response
+from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
+
+
+@pytest.fixture
+def provider(recording_server: RecordingServer, route: Route) -> RecordingServer:
+    recording_server.default_response = provider_response(route)
+    return recording_server
+
+
+@pytest.fixture
+def otel(isolate_rust_state: ExitStack) -> OtelHarness:
+    exporter: Final = InMemorySpanExporter()
+    tracer_provider: Final = TracerProvider()
+    isolate_rust_state.callback(tracer_provider.shutdown)
+    tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
+    logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=tracer_provider)
+    return OtelHarness(logger=logger, exporter=exporter)
+
+
+@pytest.fixture
+def prometheus() -> PrometheusLogger:
+    return PrometheusLogger()
+
+
+@pytest_asyncio.fixture
+async def generic_api_export(isolate_rust_state: ExitStack) -> GenericAPIExportHarness:
+    server: Final = isolate_rust_state.enter_context(recording_service())
+    server.expected_requests = None
+    logger: Final = GenericAPILogger(endpoint=f"{server.base_url}/logs", batch_size=1, log_format="single")
+    return GenericAPIExportHarness(server=server, logger=logger)

tests/test_litellm_rust/integrations/harness.py

@@ -0,0 +1,203 @@
+import asyncio
+import threading
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Final
+
+import httpx
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from prometheus_client import REGISTRY
+
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
+from litellm.integrations.gcs_bucket.gcs_bucket_base import IAM_AUTH_KEY
+from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
+from litellm.integrations.literal_ai import LiteralAILogger
+from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME, OpenTelemetry
+from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail
+from litellm.types.guardrails import GuardrailEventHooks
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer
+
+
+@dataclass(frozen=True, slots=True)
+class RunObservation:
+    call_type: str
+    model: str
+    response_cost: float
+    response_text: str
+    provider_body: Mapping[str, object]
+    rust_dispatch: bool
+
+
+@dataclass(frozen=True, slots=True)
+class RecordedPost:
+    url: str
+    headers: Mapping[str, str]
+    body: object
+
+
+class RecordingAsyncClient:
+    def __init__(
+        self, responses: tuple[Mapping[str, object], ...] = (), blocked_url_fragment: str | None = None
+    ) -> None:
+        self.posts: list[RecordedPost] = []
+        self._responses = list(responses)
+        self._blocked_url_fragment = blocked_url_fragment
+        self.accepted = threading.Event()
+        self.release = threading.Event()
+        self.release.set()
+
+    async def post(self, url: str, **kwargs: object) -> httpx.Response:
+        data: Final = kwargs.get("data")
+        json_body: Final = kwargs.get("json")
+        body: Final = json_body if json_body is not None else data
+        headers_value: Final = kwargs.get("headers")
+        headers: Final = headers_value if isinstance(headers_value, Mapping) else {}
+        self.posts.append(RecordedPost(url=url, headers=headers, body=body))
+        should_block: Final = self._blocked_url_fragment is None or self._blocked_url_fragment in url
+        if should_block:
+            self.accepted.set()
+            released: Final = await asyncio.to_thread(self.release.wait, 10)
+            if not released:
+                raise TimeoutError(f"Timed out releasing POST {url}")
+        payload: Final = self._responses.pop(0) if self._responses else {}
+        response_headers: Final = {"etag": '"test-scope"'} if "protectionScopes/compute" in url else {}
+        return httpx.Response(200, json=payload, headers=response_headers, request=httpx.Request("POST", url))
+
+
+class RecordingVertexInstance:
+    async def _ensure_access_token_async(self, **kwargs: object) -> tuple[str, str]:
+        return "test-access-token", "test-project"
+
+    def _get_token_and_url(self, **kwargs: object) -> tuple[str, str]:
+        return "test-access-token", ""
+
+
+class AsyncBoundaryLogger(CustomLogger):
+    def __init__(self, action: Callable[[], Awaitable[None]]) -> None:
+        self._action = action
+
+    async def async_log_success_event(
+        self, kwargs: object, response_obj: object, start_time: object, end_time: object
+    ) -> None:
+        await self._action()
+
+
+@dataclass(frozen=True, slots=True)
+class OtelHarness:
+    logger: OpenTelemetry
+    exporter: InMemorySpanExporter
+
+    def spans(self, name: str = LITELLM_REQUEST_SPAN_NAME) -> tuple[ReadableSpan, ...]:
+        return tuple(span for span in self.exporter.get_finished_spans() if span.name == name)
+
+    async def wait_for_spans(
+        self, name: str = LITELLM_REQUEST_SPAN_NAME, count: int = 1, timeout: float = 10
+    ) -> tuple[ReadableSpan, ...]:
+        deadline: Final = time.monotonic() + timeout
+        while len(self.spans(name)) < count:
+            if time.monotonic() >= deadline:
+                raise TimeoutError(f"Timed out waiting for {count} {name} spans; saw {self.spans(name)}")
+            await asyncio.sleep(0.01)
+        await drain_logging()
+        return self.spans(name)
+
+
+@dataclass(frozen=True, slots=True)
+class GCSLiteralAIHarness:
+    gcs: GCSBucketLogger
+    literal: LiteralAILogger
+    storage: RecordingAsyncClient
+    literal_sink: RecordingAsyncClient
+
+
+@dataclass(frozen=True, slots=True)
+class GenericAPIExportHarness:
+    server: RecordingServer
+    logger: GenericAPILogger
+
+    @property
+    def exports(self) -> tuple[RecordedRequest, ...]:
+        return tuple(request for request in self.server.requests if request.path == "/logs")
+
+
+@asynccontextmanager
+async def gcs_literalai_harness() -> AsyncIterator[GCSLiteralAIHarness]:
+    storage: Final = RecordingAsyncClient()
+    literal_sink: Final = RecordingAsyncClient()
+    gcs: Final = GCSBucketLogger(bucket_name="composition-bucket")
+    gcs.async_httpx_client = storage
+    gcs.vertex_instances[IAM_AUTH_KEY] = RecordingVertexInstance()
+    literal: Final = LiteralAILogger(literalai_api_key="test-key")
+    literal.async_httpx_client = literal_sink
+    try:
+        yield GCSLiteralAIHarness(gcs=gcs, literal=literal, storage=storage, literal_sink=literal_sink)
+    finally:
+        await gcs.aclose()
+
+
+@dataclass(frozen=True, slots=True)
+class PurviewHarness:
+    graph: RecordingAsyncClient
+    guardrail: MicrosoftPurviewDLPGuardrail
+
+
+def purview_harness(name: str) -> PurviewHarness:
+    graph: Final = RecordingAsyncClient(
+        responses=({"access_token": "test-token", "expires_in": 3600}, {}, {}, {}),
+        blocked_url_fragment="processContent",
+    )
+    graph.release.clear()
+    guardrail: Final = MicrosoftPurviewDLPGuardrail(
+        guardrail_name=name,
+        tenant_id="test-tenant",
+        client_id="test-client",
+        client_secret="test-secret",
+        event_hook=GuardrailEventHooks.logging_only,
+        default_on=True,
+    )
+    guardrail.async_handler = graph
+    return PurviewHarness(graph=graph, guardrail=guardrail)
+
+
+async def wait_for_audits(graph: RecordingAsyncClient, count: int = 2, timeout: float = 3) -> tuple[RecordedPost, ...]:
+    async with asyncio.timeout(timeout):
+        while len(tuple(post for post in graph.posts if "processContent" in post.url)) < count:
+            await asyncio.sleep(0.01)
+    return tuple(post for post in graph.posts if "processContent" in post.url)
+
+
+def metric_value(name: str, **labels: str) -> float:
+    for metric in REGISTRY.collect():
+        for sample in metric.samples:
+            if sample.name == name and all(sample.labels.get(key) == value for key, value in labels.items()):
+                return sample.value
+    return 0.0
+
+
+class ReviewGuardrail(CustomGuardrail):
+    def __init__(self, review: Callable[[object], Awaitable[object]]) -> None:
+        super().__init__(guardrail_name="rust-review", event_hook=GuardrailEventHooks.post_call, default_on=True)
+        self._review = review
+        self.call_types: list[object] = []
+
+    async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+        self.call_types.append(call_type)
+        return await self._review(response)
+
+
+class MutatingFailingLogger(CustomLogger):
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        payload: Final = kwargs["standard_logging_object"]
+        payload["metadata"]["composition_marker"] = "visible-before-failure"
+        raise RuntimeError("synthetic callback failure")
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        payload: Final = kwargs["standard_logging_object"]
+        payload["metadata"]["composition_marker"] = "visible-before-failure"
+        raise RuntimeError("synthetic callback failure")

tests/test_litellm_rust/integrations/test_composition.py

@@ -0,0 +1,548 @@
+import asyncio
+import gc
+import json
+import threading
+import weakref
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Final
+
+import pytest
+from opentelemetry.trace import StatusCode
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
+from litellm.integrations.prometheus import PrometheusLogger
+from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.conftest import Backend, isolated_backend
+from tests.test_litellm_rust.integrations import (
+    ASYNC_ROUTES,
+    DISCOVERED_ONLY_GUARDRAIL_NAMES,
+    ENTERPRISE_LOGGER_NAMES,
+    GUARDRAIL_NAMES,
+    GUARDRAIL_OBLIGATIONS,
+    LOGGER_OBLIGATIONS,
+    MESSAGES_ROUTE,
+    MESSAGES_STREAM,
+    OCR_ASYNC,
+    OCR_SYNC,
+    OSS_LOGGER_NAMES,
+    REQUIRED_GUARDRAIL_BEHAVIOR,
+    REQUIRED_LOGGER_BEHAVIOR,
+    AsyncBoundaryLogger,
+    MutatingFailingLogger,
+    OtelHarness,
+    Route,
+    RunObservation,
+    gcs_literalai_harness,
+    metric_value,
+    provider_response,
+    route_id,
+    wait_for_callback,
+)
+from tests.test_litellm_rust.support.callback_recorder import (
+    LiveReferenceLogger,
+    RecordingLogger,
+    SecondaryLiveReferenceLogger,
+    drain_logging,
+)
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
+from tests.test_litellm_rust.support.requests import (
+    CHAT_MESSAGES,
+    CHAT_MODEL,
+    CHAT_RESPONSE,
+    MESSAGES,
+    MESSAGES_EVENTS,
+)
+
+pytestmark = pytest.mark.requires_rust_extension
+
+MESSAGES_TOOLS: Final = [
+    {
+        "name": "get_weather",
+        "description": "Get current weather",
+        "input_schema": {
+            "type": "object",
+            "properties": {"city": {"type": "string"}},
+            "required": ["city"],
+        },
+    }
+]
+
+
+async def observe_backend(route: Route, backend: Backend) -> RunObservation:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = provider_response(route)
+            recorder: Final = RecordingLogger()
+            response: Final = await route.invoke(provider, callbacks=[recorder])
+            event: Final = (await wait_for_callback(route, recorder))[0]
+            payload: Final = event.kwargs["standard_logging_object"]
+            provider_body: Final = provider.requests[0].body
+            if not isinstance(provider_body, dict):
+                raise TypeError(f"Expected provider object body, got {type(provider_body).__name__}")
+            return RunObservation(
+                call_type=payload["call_type"],
+                model=payload["model"],
+                response_cost=payload["response_cost"],
+                response_text=route.response_text(response),
+                provider_body=provider_body,
+                rust_dispatch=has_rust_response_marker(response),
+            )
+
+
+def test_logger_catalogue_reconciles_with_production_registry() -> None:
+    actual_names: Final = frozenset(CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE)
+    catalogued_names: Final = frozenset(
+        name for obligation in LOGGER_OBLIGATIONS.values() for name in obligation.registration_names
+    )
+    assert OSS_LOGGER_NAMES <= actual_names
+    assert actual_names <= OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
+    assert catalogued_names == OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
+
+
+def test_guardrail_catalogue_reconciles_enum_and_runtime_discovery() -> None:
+    enum_names: Final = frozenset(integration.value for integration in SupportedGuardrailIntegrations)
+    initializer_names: Final = frozenset(guardrail_initializer_registry)
+    assert enum_names == GUARDRAIL_NAMES
+    assert initializer_names == GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES
+    assert frozenset(GUARDRAIL_OBLIGATIONS) == initializer_names
+
+
+def test_required_integrations_have_behavioral_cases() -> None:
+    logger_names: Final = frozenset(
+        name
+        for obligation in LOGGER_OBLIGATIONS.values()
+        if obligation.behavioral_cases
+        for name in obligation.registration_names
+    )
+    guardrail_names: Final = frozenset(
+        name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_cases
+    )
+    assert REQUIRED_LOGGER_BEHAVIOR <= logger_names
+    assert REQUIRED_GUARDRAIL_BEHAVIOR <= guardrail_names
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=lambda route: f"otel-prometheus-generic-api-{route.name}-success")
+async def test_export_composition(
+    route: Route, recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger
+) -> None:
+    recording_server.default_response = provider_response(route)
+    before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model)
+    recorder: Final = RecordingLogger()
+
+    with recording_service() as sink:
+        logger: Final = GenericAPILogger(endpoint=f"{sink.base_url}/logs", batch_size=1, log_format="single")
+        await route.invoke(recording_server, callbacks=[otel.logger, prometheus, logger, recorder])
+        await wait_for_callback(route, recorder)
+
+        exports: Final = tuple(request for request in sink.requests if request.path == "/logs")
+        spans: Final = await otel.wait_for_spans()
+        assert len(exports) == 1
+        assert exports[0].body["call_type"] == route.call_type
+        assert exports[0].body["response_cost"] == pytest.approx(route.expected_cost)
+        assert len(spans) == 1
+        assert spans[0].status.status_code is StatusCode.OK
+        assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE), ids=route_id)
+async def test_public_sdk_python_rust_composition_parity(route: Route) -> None:
+    python: Final = await observe_backend(route, "python")
+    rust: Final = await observe_backend(route, "rust")
+    assert python.call_type == rust.call_type == route.call_type
+    assert python.model == rust.model == route.provider_model
+    assert python.response_cost == pytest.approx(rust.response_cost)
+    assert python.response_text == rust.response_text == route.expected_text
+    python_body: Final = {**python.provider_body, "stream": python.provider_body.get("stream", False)}
+    rust_body: Final = {**rust.provider_body, "stream": rust.provider_body.get("stream", False)}
+    assert python_body == rust_body
+    assert python.rust_dispatch is False
+    assert rust.rust_dispatch is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_terminal_callbacks_share_live_objects_within_one_run(backend: Backend) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = provider_response(MESSAGES_ROUTE)
+            first: Final = LiveReferenceLogger()
+            second: Final = SecondaryLiveReferenceLogger()
+            response: Final = await MESSAGES_ROUTE.invoke(provider, callbacks=[first, second])
+            first_event: Final = (await first.wait_for_async())[0]
+            second_event: Final = (await second.wait_for_async())[0]
+            assert first_event.kwargs is second_event.kwargs
+            assert first_event.response is second_event.response
+            assert has_rust_response_marker(response) is (backend == "rust")
+            first.release()
+            second.release()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_per_logger_redaction_preserves_the_shared_unredacted_payload(backend: Backend) -> None:
+    class Capture(CustomLogger):
+        def __init__(self, redact: bool) -> None:
+            super().__init__(turn_off_message_logging=redact)
+            self.kwargs: object = None
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            self.kwargs = kwargs
+
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = provider_response(MESSAGES_ROUTE)
+            redacted: Final = Capture(True)
+            plain: Final = Capture(False)
+            await MESSAGES_ROUTE.invoke(provider, callbacks=[redacted, plain])
+            await drain_logging()
+
+            redacted_kwargs: Final = redacted.kwargs
+            plain_kwargs: Final = plain.kwargs
+            assert isinstance(redacted_kwargs, dict)
+            assert isinstance(plain_kwargs, dict)
+            redacted_payload: Final = redacted_kwargs["standard_logging_object"]
+            plain_payload: Final = plain_kwargs["standard_logging_object"]
+            assert redacted_kwargs is not plain_kwargs
+            assert redacted_payload is not plain_payload
+            assert redacted_payload["metadata"] is plain_payload["metadata"]
+            assert "redacted-by-litellm" in json.dumps(redacted_payload["messages"])
+            assert MESSAGES_ROUTE.expected_text in json.dumps(plain_payload["response"])
+            assert "redacted-by-litellm" not in json.dumps(plain_payload)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected"))
+async def test_terminal_respects_proxy_deferred_completion_gate(backend: Backend, accepted: bool) -> None:
+    from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = provider_response(MESSAGES_ROUTE)
+            recorder: Final = RecordingLogger()
+            logger: Final = Logging(
+                model=MESSAGES_ROUTE.provider_model,
+                messages=MESSAGES,
+                stream=False,
+                call_type=MESSAGES_ROUTE.call_type,
+                start_time=datetime.now(),
+                litellm_call_id="deferred-completion",
+                function_id="deferred-completion",
+                dynamic_async_success_callbacks=[recorder],
+                dynamic_async_failure_callbacks=[recorder],
+            )
+            logger._defer_async_logging = True
+            response: Final = await MESSAGES_ROUTE.invoke(provider, litellm_logging_obj=logger)
+            await asyncio.sleep(0)
+            assert "async_log_success_event" not in recorder.names
+
+            ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
+                logging_obj=logger,
+                exception_raised=not accepted,
+            )
+            if accepted:
+                accepted_events: Final = await recorder.wait_for_async("async_log_success_event")
+                assert len(accepted_events) == 1
+                assert "async_log_failure_event" not in recorder.names
+            else:
+                error: Final = RuntimeError("post-call guardrail rejected response")
+                await logger.async_failure_handler(error, str(error))
+                rejected_events: Final = await recorder.wait_for_async("async_log_failure_event")
+                assert len(rejected_events) == 1
+                assert "async_log_success_event" not in recorder.names
+            assert has_rust_response_marker(response) is (backend == "rust")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_suspended_terminal_callback_owns_payload_until_it_finishes(backend: Backend) -> None:
+    class Root:
+        pass
+
+    class Suspended(CustomLogger):
+        def __init__(self) -> None:
+            super().__init__()
+            self.started = asyncio.Event()
+            self.release = asyncio.Event()
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            assert kwargs["litellm_params"]["metadata"]["retained"] is not None
+            self.started.set()
+            await self.release.wait()
+
+    async def invoke() -> weakref.ReferenceType[Root]:
+        async with isolated_backend(backend):
+            with recording_service() as provider:
+                provider.default_response = provider_response(MESSAGES_ROUTE)
+                callback: Final = Suspended()
+                root = Root()
+                reference: Final = weakref.ref(root)
+                response: Final = await MESSAGES_ROUTE.invoke(
+                    provider,
+                    callbacks=[callback],
+                    metadata={"retained": root},
+                )
+                del root
+                await asyncio.wait_for(callback.started.wait(), timeout=2)
+                gc.collect()
+                assert reference() is not None
+                callback.release.set()
+                await drain_logging()
+                del response
+                return reference
+
+    reference: Final = await invoke()
+    await asyncio.sleep(0)
+    gc.collect()
+    assert reference() is None
+
+
+async def observe_retry(backend: Backend) -> tuple[tuple[str, ...], bool]:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.expected_requests = 2
+            provider.enqueue(ResponseSpec(body={"message": "retry this attempt"}, status=500))
+            provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
+            recorder: Final = RecordingLogger()
+            response: Final = await litellm.acompletion(
+                model=CHAT_MODEL,
+                messages=CHAT_MESSAGES,
+                api_key="test-key",
+                api_base=provider.base_url,
+                callbacks=[recorder],
+                num_retries=1,
+            )
+            await drain_logging()
+            return recorder.names, has_rust_response_marker(response)
+
+
+@pytest.mark.asyncio
+async def test_retry_attempt_callback_sequence_matches_python() -> None:
+    python_names, python_rust_dispatch = await observe_retry("python")
+    rust_names, rust_rust_dispatch = await observe_retry("rust")
+    assert rust_names == python_names
+    assert rust_names.count("log_pre_api_call") == 2
+    assert rust_names.count("async_log_failure_event") == 1
+    assert rust_names.count("async_log_success_event") == 0
+    assert python_rust_dispatch is False
+    assert rust_rust_dispatch is True
+
+
+@pytest.mark.asyncio
+async def test_failing_callback_preserves_prior_mutation_and_later_exporters() -> None:
+    with recording_service() as provider, recording_service() as sink:
+        provider.default_response = provider_response(OCR_ASYNC)
+        logger: Final = GenericAPILogger(endpoint=f"{sink.base_url}/logs", batch_size=1, log_format="single")
+        recorder: Final = RecordingLogger()
+        response: Final = await OCR_ASYNC.invoke(provider, callbacks=[MutatingFailingLogger(), logger, recorder])
+        events: Final = await wait_for_callback(OCR_ASYNC, recorder)
+        exports: Final = tuple(request for request in sink.requests if request.path == "/logs")
+        assert response.pages[0].markdown == OCR_ASYNC.expected_text
+        assert events[0].kwargs["standard_logging_object"]["metadata"]["composition_marker"] == "visible-before-failure"
+        assert len(exports) == 1
+        assert exports[0].body["metadata"]["composition_marker"] == "visible-before-failure"
+
+
+@dataclass(frozen=True, slots=True)
+class SerializationSchedule:
+    callback_order: tuple[str, str]
+    flush_immediately: bool
+    gcs_has_tools: bool
+
+
+SERIALIZATION_SCHEDULES: Final = (
+    pytest.param(
+        SerializationSchedule(("gcs", "literalai"), True, True), id="gcs-serializes-before-literalai-mutation"
+    ),
+    pytest.param(
+        SerializationSchedule(("gcs", "literalai"), False, False), id="gcs-serializes-after-literalai-mutation"
+    ),
+    pytest.param(SerializationSchedule(("literalai", "gcs"), False, False), id="literalai-mutates-before-gcs-enqueue"),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+@pytest.mark.parametrize("schedule", SERIALIZATION_SCHEDULES)
+async def test_gcs_literalai_serialization_schedule(
+    backend: Backend, schedule: SerializationSchedule, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from litellm.proxy import proxy_server
+
+    monkeypatch.setattr(proxy_server, "premium_user", True)
+    monkeypatch.setenv("GCS_BATCH_SIZE", "1" if schedule.flush_immediately else "100")
+    monkeypatch.setenv("GCS_BUCKET_NAME", "composition-bucket")
+    monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600")
+    monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "true")
+    monkeypatch.setenv("LITERAL_BATCH_SIZE", "1")
+
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = provider_response(MESSAGES_ROUTE)
+            async with gcs_literalai_harness() as harness:
+                callbacks_by_name: Final = {"gcs": harness.gcs, "literalai": harness.literal}
+                ordered_callbacks: Final = [callbacks_by_name[name] for name in schedule.callback_order]
+                callbacks: Final = (
+                    [harness.gcs, AsyncBoundaryLogger(harness.gcs.flush_queue), harness.literal]
+                    if schedule.flush_immediately
+                    else ordered_callbacks
+                )
+                recorder: Final = RecordingLogger()
+                response: Final = await MESSAGES_ROUTE.invoke(
+                    provider, tools=MESSAGES_TOOLS, callbacks=[*callbacks, recorder]
+                )
+                await wait_for_callback(MESSAGES_ROUTE, recorder)
+                if not schedule.flush_immediately:
+                    await harness.gcs.flush_queue()
+
+                gcs_payload_value: Final = harness.storage.posts[0].body
+                if not isinstance(gcs_payload_value, str):
+                    raise TypeError(f"Expected serialized GCS payload, got {type(gcs_payload_value).__name__}")
+                gcs_payload: Final = json.loads(gcs_payload_value)
+                literal_body: Final = harness.literal_sink.posts[0].body
+                if not isinstance(literal_body, dict):
+                    raise TypeError(f"Expected LiteralAI request body, got {type(literal_body).__name__}")
+                generation: Final = literal_body["variables"]["generation_0"]
+                assert ("tools" in gcs_payload["model_parameters"]) is schedule.gcs_has_tools
+                assert generation["tools"] == MESSAGES_TOOLS
+                assert len(harness.storage.posts) == 1
+                assert len(harness.literal_sink.posts) == 1
+                assert has_rust_response_marker(response) is (backend == "rust")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_interrupted_stream_emits_terminal_only_when_consumer_closes(backend: Backend, otel: OtelHarness) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
+            stream: Final = await MESSAGES_STREAM.open_stream(provider, callbacks=[otel.logger])
+            first_chunk: Final = await anext(stream)
+            await drain_logging()
+            assert otel.spans() == ()
+            assert has_rust_response_marker(stream) is (backend == "rust")
+            await stream.aclose()
+            await drain_logging()
+            assert first_chunk is not None
+            assert len(otel.spans()) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+@pytest.mark.parametrize(
+    ("route", "phase"),
+    (
+        (OCR_ASYNC, "pre"),
+        (MESSAGES_ROUTE, "pre"),
+        (MESSAGES_STREAM, "pre"),
+        (OCR_ASYNC, "success"),
+        (MESSAGES_ROUTE, "success"),
+    ),
+    ids=("ocr-pre", "messages-pre", "messages-stream-pre", "ocr-success", "messages-success"),
+)
+async def test_deployment_rejection_logs_failure_without_deployment_failure(
+    backend: Backend, route: Route, phase: str
+) -> None:
+    class Reject(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: CallTypes | None
+        ) -> dict[str, object]:
+            if phase == "pre":
+                raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
+            return kwargs
+
+        async def async_post_call_success_deployment_hook(
+            self, request_data: Mapping[str, object], response: object, call_type: CallTypes | None
+        ) -> object:
+            raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
+
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.expected_requests = 0 if phase == "pre" else 1
+            provider.default_response = provider_response(route)
+            recorder: Final = RecordingLogger()
+            litellm.callbacks.extend((Reject(), recorder))
+
+            with pytest.raises(litellm.BadRequestError, match="deployment rejected"):
+                await route.invoke(provider)
+            await recorder.wait_for_async("async_log_failure_event")
+
+            assert len(provider.requests) == provider.expected_requests
+            assert "async_log_success_event" not in recorder.names
+            assert recorder.names.count("async_post_call_failure_deployment_hook") == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("rust",))
+@pytest.mark.parametrize("ending", ("provider_error", "truncated", "close"))
+async def test_established_stream_completion_does_not_notify_deployment_failure(backend: Backend, ending: str) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = ResponseSpec(
+                body=None,
+                events=(
+                    (
+                        *MESSAGES_EVENTS[:1],
+                        (
+                            "error",
+                            {"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}},
+                        ),
+                    )
+                    if ending == "provider_error"
+                    else MESSAGES_EVENTS[:-1]
+                    if ending == "truncated"
+                    else MESSAGES_EVENTS
+                ),
+            )
+            release: Final = threading.Event()
+            if ending == "close":
+                provider.enqueue(
+                    ResponseSpec(
+                        body=None,
+                        chunks=tuple(
+                            f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS
+                        ),
+                        release=release,
+                    )
+                )
+            recorder: Final = RecordingLogger()
+            litellm.callbacks.append(recorder)
+            try:
+                stream: Final = await MESSAGES_STREAM.open_stream(provider)
+                assert has_rust_response_marker(stream) is (backend == "rust")
+                assert recorder.names.count("async_pre_call_deployment_hook") == 1
+                assert "async_post_call_failure_deployment_hook" not in recorder.names
+
+                if ending == "close":
+                    assert await anext(stream)
+                    await asyncio.wait_for(stream.aclose(), timeout=1)
+                    await stream.aclose()
+                    assert "async_log_success_event" not in recorder.names
+                    assert "async_log_failure_event" not in recorder.names
+                    release.set()
+                    await recorder.wait_for_async("async_log_success_event")
+                else:
+                    try:
+                        async for _ in stream:
+                            pass
+                    except litellm.APIError:
+                        pass
+                await drain_logging()
+            finally:
+                release.set()
+
+            assert len(provider.requests) == 1
+            assert recorder.names.count("async_log_failure_event") == (0 if ending == "close" else 1), recorder.names
+            assert recorder.names.count("async_log_success_event") == (1 if ending == "close" else 0), recorder.names
+            assert recorder.names.count("async_post_call_failure_deployment_hook") == 0

tests/test_litellm_rust/integrations/test_exporters.py

@@ -0,0 +1,223 @@
+import json
+from typing import Final
+
+import pytest
+from opentelemetry.trace import StatusCode
+from prometheus_client import CollectorRegistry, Counter
+
+import litellm
+from litellm.integrations.prometheus import PrometheusLogger
+from litellm.litellm_core_utils import litellm_logging
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry
+from tests.test_litellm_rust.support.requests import MESSAGES_EVENTS
+from tests.test_litellm_rust.integrations import (
+    ALL_ROUTES,
+    ASYNC_ROUTES,
+    MESSAGES_ROUTE,
+    MESSAGES_STREAM,
+    NON_STREAM_ASYNC_ROUTES,
+    OCR_ASYNC,
+    OCR_SYNC,
+    GenericAPIExportHarness,
+    OtelHarness,
+    Route,
+    metric_value,
+    route_id,
+    wait_for_callback,
+)
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+FAILURE_RESPONSE: Final = ResponseSpec(body={"message": "provider unavailable"}, status=500)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id)
+async def test_generic_api_logger_exports_success_over_http(
+    route: Route,
+    provider: RecordingServer,
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    recorder: Final = RecordingLogger()
+    await route.invoke(provider, callbacks=[generic_api_export.logger, recorder])
+    await wait_for_callback(route, recorder)
+
+    assert len(generic_api_export.exports) == 1
+    payload: Final = generic_api_export.exports[0].body
+    assert payload["status"] == "success"
+    assert payload["call_type"] == route.call_type
+    assert payload["model"] == route.provider_model
+    assert payload["response_cost"] == pytest.approx(route.expected_cost)
+    assert route.expected_text in json.dumps(payload["response"])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
+async def test_generic_api_logger_exports_provider_failure_over_http(
+    route: Route,
+    provider: RecordingServer,
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    provider.enqueue(FAILURE_RESPONSE)
+    recorder: Final = RecordingLogger()
+
+    try:
+        with pytest.raises(litellm.InternalServerError):
+            await route.invoke(provider, callbacks=[generic_api_export.logger, recorder], num_retries=0)
+    finally:
+        await drain_logging()
+    await wait_for_callback(route, recorder, outcome="failure")
+
+    assert len(provider.requests) == 1
+    assert len(generic_api_export.exports) == 1
+    payload: Final = generic_api_export.exports[0].body
+    assert payload["status"] == "failure"
+    assert payload["call_type"] == route.call_type
+    assert payload["error_information"]["error_class"] == "InternalServerError"
+    assert payload["error_information"]["error_code"] == "500"
+
+
+def test_prometheus_registry_restores_collectors_after_failure() -> None:
+    registry: Final = CollectorRegistry()
+    original: Final = Counter("original", "Original collector", registry=registry)
+    original.inc(2)
+
+    def failing_test() -> None:
+        with isolated_prometheus_registry(registry):
+            assert registry.get_sample_value("original_total") is None
+            Counter("original", "Temporary replacement", registry=registry).inc(7)
+            Counter("temporary", "Temporary collector", registry=registry).inc()
+            raise RuntimeError("test failed")
+
+    with pytest.raises(RuntimeError, match="test failed"):
+        failing_test()
+    assert registry.get_sample_value("original_total") == 2
+    assert registry.get_sample_value("temporary_total") is None
+    registry.unregister(original)
+    assert registry.get_sample_value("original_total") is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id)
+async def test_otel_emits_one_request_span_on_success(
+    route: Route, provider: RecordingServer, otel: OtelHarness
+) -> None:
+    await route.invoke(provider, callbacks=[otel.logger])
+    spans: Final = await otel.wait_for_spans()
+    assert len(spans) == 1
+    span: Final = spans[0]
+    assert span.status.status_code is StatusCode.OK
+    assert span.attributes["llm.request.type"] == route.call_type
+    assert span.attributes["gen_ai.request.model"] == route.provider_model
+    assert json.loads(span.attributes["hidden_params"])["response_cost"] == pytest.approx(route.expected_cost)
+
+
+@pytest.mark.asyncio
+async def test_otel_stream_span_appears_only_after_exhaustion(
+    otel: OtelHarness, recording_server: RecordingServer
+) -> None:
+    recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
+    stream: Final = await MESSAGES_STREAM.open_stream(recording_server, callbacks=[otel.logger])
+    await drain_logging()
+    assert otel.spans() == ()
+    chunks: Final = [chunk async for chunk in stream]
+    assert chunks
+    assert len(await otel.wait_for_spans()) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC), ids=route_id)
+async def test_otel_emits_one_error_span_on_provider_failure(
+    route: Route, provider: RecordingServer, otel: OtelHarness
+) -> None:
+    provider.enqueue(FAILURE_RESPONSE)
+    with pytest.raises(litellm.InternalServerError):
+        await route.invoke(provider, callbacks=[otel.logger])
+    spans: Final = await otel.wait_for_spans()
+    assert len(spans) == 1
+    assert spans[0].status.status_code is StatusCode.ERROR
+    exception_events: Final = [event for event in spans[0].events if event.name == "exception"]
+    assert len(exception_events) == 1
+    assert "InternalServerError" in exception_events[0].attributes["exception.type"]
+    assert spans[0].attributes["llm.request.type"] == route.call_type
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id)
+async def test_otel_and_custom_logger_observe_same_standard_logging_object(
+    route: Route, provider: RecordingServer, otel: OtelHarness
+) -> None:
+    recorder: Final = RecordingLogger()
+    await route.invoke(provider, callbacks=[otel.logger, recorder])
+    payload: Final = (await wait_for_callback(route, recorder))[0].kwargs["standard_logging_object"]
+    span: Final = (await otel.wait_for_spans())[0]
+    assert payload["call_type"] == route.call_type
+    assert payload["model"] == route.provider_model
+    assert json.loads(span.attributes["hidden_params"]) == payload["hidden_params"]
+    assert span.attributes["llm.request.type"] == payload["call_type"]
+    assert span.attributes["litellm.provider.model"] == payload["model"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id)
+async def test_prometheus_counts_one_successful_request(
+    route: Route, provider: RecordingServer, prometheus: PrometheusLogger
+) -> None:
+    before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model)
+    recorder: Final = RecordingLogger()
+    await route.invoke(provider, callbacks=[prometheus, recorder])
+    await wait_for_callback(route, recorder)
+    assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1
+    assert metric_value("litellm_llm_api_failed_requests_metric_total", model=route.provider_model) == 0
+
+
+@pytest.mark.asyncio
+async def test_prometheus_counts_tokens_from_messages_usage(
+    recording_server: RecordingServer, prometheus: PrometheusLogger
+) -> None:
+    recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response)
+    recorder: Final = RecordingLogger()
+    await MESSAGES_ROUTE.invoke(recording_server, callbacks=[prometheus, recorder])
+    await wait_for_callback(MESSAGES_ROUTE, recorder)
+    assert metric_value("litellm_input_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 5
+    assert metric_value("litellm_output_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 4
+
+
+@pytest.mark.asyncio
+async def test_prometheus_counts_one_failed_request(
+    otel: OtelHarness, prometheus: PrometheusLogger, recording_server: RecordingServer
+) -> None:
+    recording_server.enqueue(FAILURE_RESPONSE)
+    with pytest.raises(litellm.InternalServerError):
+        await OCR_ASYNC.invoke(recording_server, callbacks=[prometheus, otel.logger])
+    assert len(await otel.wait_for_spans()) == 1
+    assert metric_value("litellm_llm_api_failed_requests_metric_total", model=OCR_ASYNC.provider_model) == 1
+    assert metric_value("litellm_requests_metric_total", model=OCR_ASYNC.provider_model) == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
+async def test_prometheus_by_string_name_is_initialized_once(route: Route, provider: RecordingServer) -> None:
+    provider.expected_requests = 2
+    litellm.success_callback = ["prometheus"]  # test-quality-ok: public registration; fixture restores globals
+    recorder: Final = RecordingLogger()
+    await route.invoke(provider, callbacks=[recorder])
+    await route.invoke(provider, callbacks=[recorder])
+    await wait_for_callback(route, recorder, count=2)
+    instances: Final = [cb for cb in litellm_logging._in_memory_loggers if isinstance(cb, PrometheusLogger)]  # pyright: ignore[reportPrivateUsage]  # string-name cache has no public accessor
+    assert len(instances) == 1
+    assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 2
+    assert "prometheus" not in litellm.success_callback
+    assert instances[0] in litellm._async_success_callback  # pyright: ignore[reportPrivateUsage]  # callback registry has no public accessor
+
+
+@pytest.mark.asyncio
+async def test_sync_ocr_reaches_sync_hooks_only(
+    recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger
+) -> None:
+    recording_server.default_response = ResponseSpec(body=OCR_SYNC.provider_response)
+    await OCR_SYNC.invoke(recording_server, callbacks=[otel.logger, prometheus])
+    assert len(await otel.wait_for_spans()) == 1
+    assert metric_value("litellm_requests_metric_total", model=OCR_SYNC.provider_model) == 0

tests/test_litellm_rust/integrations/test_guardrails.py

@@ -0,0 +1,697 @@
+import asyncio
+import gc
+import json
+import threading
+import weakref
+from collections.abc import AsyncIterator
+from datetime import datetime
+from typing import Final
+from unittest.mock import MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+import litellm
+from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
+from litellm.integrations.literal_ai import LiteralAILogger
+from litellm.integrations.rubrik import RubrikLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import AzureContentSafetyTextModerationGuardrail
+from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler
+from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
+from litellm.proxy.utils import ProxyLogging
+from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+from tests.test_litellm_rust.conftest import Backend, isolated_backend
+from tests.test_litellm_rust.support.requests import CHAT_MESSAGES, CHAT_MODEL, CHAT_RESPONSE, MESSAGES
+from tests.test_litellm_rust.integrations import (
+    MESSAGES_ROUTE,
+    MESSAGES_STREAM,
+    NON_STREAM_ASYNC_ROUTES,
+    GenericAPIExportHarness,
+    OtelHarness,
+    ReviewGuardrail,
+    Route,
+    provider_response,
+    purview_harness,
+    route_id,
+    wait_for_audits,
+    wait_for_callback,
+)
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
+
+pytestmark = pytest.mark.requires_rust_extension
+
+AZURE_MODERATION_ALLOW_RESPONSE: Final = {
+    "blocklistsMatch": [],
+    "categoriesAnalysis": [
+        {"category": "Hate", "severity": 0},
+        {"category": "Sexual", "severity": 0},
+        {"category": "SelfHarm", "severity": 0},
+        {"category": "Violence", "severity": 0},
+    ],
+}
+AZURE_MODERATION_BLOCK_RESPONSE: Final = {
+    **AZURE_MODERATION_ALLOW_RESPONSE,
+    "categoriesAnalysis": [{"category": "Violence", "severity": 6}],
+}
+STREAM_SECRET: Final = "078-05-1120"
+STREAM_MASK: Final = "[MASKED]"
+
+
+class BufferedMaskingGuardrail(CustomGuardrail):
+    def __init__(self) -> None:
+        super().__init__(guardrail_name="stream-mask", event_hook=GuardrailEventHooks.post_call, default_on=True)
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterator[object],
+        request_data: dict,
+    ) -> AsyncIterator[bytes]:
+        chunks: Final = tuple([chunk async for chunk in response])
+        body: Final = b"".join(chunk for chunk in chunks if isinstance(chunk, bytes))
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response={"masked": True},
+            request_data=request_data,
+            guardrail_status="success",
+            masked_entity_count={"TEST_IDENTIFIER": 1},
+        )
+        yield body.replace(STREAM_SECRET.encode(), STREAM_MASK.encode())
+
+
+class CleanupFailureGuardrail(CustomGuardrail):
+    def __init__(self) -> None:
+        super().__init__(guardrail_name="cleanup-failure", event_hook=GuardrailEventHooks.post_call, default_on=True)
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterator[object],
+        request_data: dict,
+    ) -> AsyncIterator[object]:
+        try:
+            async for chunk in response:
+                yield chunk
+        finally:
+            self.add_standard_logging_guardrail_information_to_request_data(
+                guardrail_json_response="guardrail cleanup failed",
+                request_data=request_data,
+                guardrail_status="guardrail_failed_to_respond",
+            )
+            raise RuntimeError("guardrail cleanup failed")
+
+
+class FailingMaskingGuardrail(CustomGuardrail):
+    def __init__(self) -> None:
+        super().__init__(guardrail_name="mask-failure", event_hook=GuardrailEventHooks.post_call, default_on=True)
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncIterator[object],
+        request_data: dict,
+    ) -> AsyncIterator[object]:
+        tuple([chunk async for chunk in response])
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response="stream masking failed",
+            request_data=request_data,
+            guardrail_status="guardrail_failed_to_respond",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail={"error": "stream masking failed", "guardrail": self.guardrail_name},
+        )
+        yield
+
+
+def stream_events(text: str) -> tuple[tuple[str, object], ...]:
+    return tuple(
+        (event, {**data, "delta": {**data["delta"], "text": text}})
+        if event == "content_block_delta" and isinstance(data, dict)
+        else (event, data)
+        for event, data in provider_response(MESSAGES_STREAM).events
+    )
+
+
+def streaming_logger(recorder: RecordingLogger, callbacks: list[object] | None = None) -> Logging:
+    terminal_callbacks: Final = [recorder, *(callbacks or [])]
+    return Logging(
+        model=MESSAGES_STREAM.provider_model,
+        messages=[{"role": "user", "content": "mask the response"}],
+        stream=True,
+        call_type="anthropic_messages",
+        start_time=datetime.now(),
+        litellm_call_id="stream-mask",
+        function_id="stream-mask",
+        dynamic_async_success_callbacks=terminal_callbacks,
+        dynamic_async_failure_callbacks=terminal_callbacks,
+    )
+
+
+def arm_guarded_stream(stream: AsyncIterator[object], request_data: dict, logger: Logging) -> AsyncIterator[object]:
+    processor: Final = object.__new__(ProxyBaseLLMRequestProcessing)
+    processor.data = request_data
+    user: Final = UserAPIKeyAuth(request_route="/v1/messages")
+    processor._arm_deferred_stream_dispatch(stream, "anthropic_messages", user, logger)
+    return ProxyLogging(user_api_key_cache=MagicMock()).async_post_call_streaming_iterator_hook(
+        response=stream,
+        user_api_key_dict=user,
+        request_data=request_data,
+    )
+
+
+def azure_text_moderation(server: RecordingServer) -> AzureContentSafetyTextModerationGuardrail:
+    return AzureContentSafetyTextModerationGuardrail(
+        guardrail_name="azure-text-review",
+        api_key="test-azure-key",
+        api_base=server.base_url,
+        event_hook=GuardrailEventHooks.post_call,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_stream_masking_reaches_client_loggers_and_exporters(
+    backend: Backend,
+    otel: OtelHarness,
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            wire: Final = b"".join(
+                f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in stream_events(STREAM_SECRET)
+            )
+            split: Final = wire.index(STREAM_SECRET.encode()) + 4
+            provider.default_response = ResponseSpec(body=None, chunks=(wire[:split], wire[split:]))
+            guardrail: Final = BufferedMaskingGuardrail()
+            recorder: Final = RecordingLogger()
+            logger: Final = streaming_logger(recorder, [otel.logger, generic_api_export.logger])
+            request_data: Final = {
+                "model": MESSAGES_STREAM.provider_model,
+                "messages": MESSAGES,
+                "guardrails": [guardrail.guardrail_name],
+                "metadata": {},
+                "litellm_logging_obj": logger,
+            }
+            litellm.callbacks.append(guardrail)
+            source: Final = await MESSAGES_STREAM.open_stream(
+                provider,
+                litellm_logging_obj=logger,
+                guardrails=[guardrail.guardrail_name],
+            )
+            guarded: Final = arm_guarded_stream(source, request_data, logger)
+
+            assert "async_log_success_event" not in recorder.names
+            payload: Final = b"".join([chunk async for chunk in guarded if isinstance(chunk, bytes)])
+            events: Final = await recorder.wait_for_async("async_log_success_event")
+            spans: Final = await otel.wait_for_spans()
+
+            assert STREAM_MASK.encode() in payload
+            assert STREAM_SECRET.encode() not in payload
+            assert payload.count(b"event: message_stop") == 1
+            assert len(events) == 1
+            assert events[0].response.choices[0].message.content == STREAM_MASK
+            assert STREAM_SECRET not in json.dumps(events[0].kwargs, default=str)
+            assert STREAM_SECRET not in json.dumps(dict(spans[0].attributes), default=str)
+            assert STREAM_MASK in json.dumps(dict(spans[0].attributes), default=str)
+            async with asyncio.timeout(10):
+                while not generic_api_export.exports:
+                    await asyncio.sleep(0.01)
+            exported: Final = json.dumps(generic_api_export.exports[0].body, default=str)
+            assert STREAM_SECRET not in exported
+            assert STREAM_MASK in exported
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_upstream_stream_failure_is_not_masked_by_guardrail_cleanup(
+    backend: Backend,
+) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = ResponseSpec(
+                body=None,
+                events=(
+                    provider_response(MESSAGES_STREAM).events[0],
+                    (
+                        "error",
+                        {
+                            "type": "error",
+                            "error": {"type": "overloaded_error", "message": "upstream overloaded"},
+                        },
+                    ),
+                ),
+            )
+            guardrail: Final = CleanupFailureGuardrail()
+            recorder: Final = RecordingLogger()
+            logger: Final = streaming_logger(recorder)
+            request_data: Final = {
+                "model": MESSAGES_STREAM.provider_model,
+                "messages": MESSAGES,
+                "guardrails": [guardrail.guardrail_name],
+                "metadata": {},
+                "litellm_logging_obj": logger,
+            }
+            litellm.callbacks.append(guardrail)
+            source: Final = await MESSAGES_STREAM.open_stream(
+                provider,
+                litellm_logging_obj=logger,
+                guardrails=[guardrail.guardrail_name],
+            )
+            guarded: Final = arm_guarded_stream(source, request_data, logger)
+
+            with pytest.raises(litellm.APIError) as caught:
+                async for _ in guarded:
+                    pass
+            guardrail_info: Final = request_data["metadata"]["standard_logging_guardrail_information"]
+
+            assert "guardrail cleanup failed" not in str(caught.value)
+            assert request_data["metadata"]["stream_guardrail_cleanup_error"] == {
+                "type": "RuntimeError",
+                "message": "guardrail cleanup failed",
+            }
+            if backend == "rust":
+                events: Final = await recorder.wait_for_async("async_log_failure_event")
+                assert len(events) == 1
+                assert "guardrail cleanup failed" not in str(events[0].kwargs["exception"])
+            else:
+                assert "async_log_success_event" not in recorder.names
+            assert guardrail_info[-1]["guardrail_status"] == "guardrail_failed_to_respond"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_stream_masking_failure_uses_guardrail_error_and_logs_once(backend: Backend) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = ResponseSpec(body=None, events=stream_events(STREAM_SECRET))
+            guardrail: Final = FailingMaskingGuardrail()
+            recorder: Final = RecordingLogger()
+            logger: Final = streaming_logger(recorder)
+            request_data: Final = {
+                "model": MESSAGES_STREAM.provider_model,
+                "messages": MESSAGES,
+                "guardrails": [guardrail.guardrail_name],
+                "metadata": {},
+                "litellm_logging_obj": logger,
+            }
+            litellm.callbacks.append(guardrail)
+            source: Final = await MESSAGES_STREAM.open_stream(
+                provider,
+                litellm_logging_obj=logger,
+                guardrails=[guardrail.guardrail_name],
+            )
+
+            with pytest.raises(HTTPException) as caught:
+                async for _ in arm_guarded_stream(source, request_data, logger):
+                    pass
+            await logger.async_failure_handler(caught.value, "stream masking failed")
+            events: Final = await recorder.wait_for_async("async_log_failure_event")
+
+            assert caught.value.status_code == 500
+            assert caught.value.detail["error"] == "stream masking failed"
+            assert caught.value.detail["guardrail"] == guardrail.guardrail_name
+            assert len(events) == 1
+            assert events[0].kwargs["exception"] is caught.value
+            assert "async_log_success_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_guarded_stream_cancel_drains_and_releases_roots(backend: Backend) -> None:
+    class NamesOnlyLogger(RecordingLogger):
+        def _record(self, name: str, kwargs: object = None, response: object = None) -> None:
+            super()._record(name)
+
+    class Root:
+        pass
+
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            release: Final = threading.Event()
+            wire: Final = tuple(
+                f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in stream_events(STREAM_SECRET)
+            )
+            provider.enqueue(ResponseSpec(body=None, chunks=wire, release=release))
+            guardrail: Final = BufferedMaskingGuardrail()
+            recorder: Final = NamesOnlyLogger()
+            logger: Final = streaming_logger(recorder)
+            root = Root()
+            reference: Final = weakref.ref(root)
+            request_data: Final = {
+                "model": MESSAGES_STREAM.provider_model,
+                "messages": MESSAGES,
+                "guardrails": [guardrail.guardrail_name],
+                "metadata": {},
+                "litellm_logging_obj": logger,
+            }
+            litellm.callbacks.append(guardrail)
+            source: Final = await MESSAGES_STREAM.open_stream(
+                provider,
+                litellm_logging_obj=logger,
+                guardrails=[guardrail.guardrail_name],
+                metadata={"retained": root},
+            )
+            first_chunk: Final = asyncio.Event()
+
+            async def observed_source() -> AsyncIterator[bytes]:
+                async for chunk in source:
+                    first_chunk.set()
+                    yield chunk
+
+            guarded = arm_guarded_stream(observed_source(), request_data, logger)
+            del root
+            task: Final = asyncio.create_task(anext(guarded))
+            await provider.wait_for_requests(1)
+            await asyncio.wait_for(first_chunk.wait(), 5)
+            task.cancel()
+
+            try:
+                with pytest.raises(asyncio.CancelledError):
+                    await task
+                await guarded.aclose()
+                release.set()
+                events: Final = await recorder.wait_for_async("async_log_success_event")
+                del guarded
+                del source
+                gc.collect()
+
+                assert len(events) == 1
+                assert "async_log_failure_event" not in recorder.names
+                del events
+                del task
+                del request_data
+                del logger
+                gc.collect()
+                assert reference() is None
+            finally:
+                release.set()
+
+
+@pytest.mark.asyncio
+async def test_crowdstrike_redaction_reaches_native_chat_provider_and_exporter(
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    with recording_service() as provider, recording_service() as guard:
+        provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
+        guard.default_response = ResponseSpec(
+            body={
+                "result": {
+                    "blocked": False,
+                    "transformed": True,
+                    "guard_output": {
+                        "messages": [
+                            {"role": "system", "content": "Keep the answer short"},
+                            {"role": "user", "content": "Employee SSN: <US_SSN>"},
+                        ]
+                    },
+                }
+            }
+        )
+        guardrail: Final = CrowdStrikeAIDRHandler(
+            guardrail_name="crowdstrike-redaction",
+            api_key="test-crowdstrike-key",
+            api_base=guard.base_url,
+            event_hook=GuardrailEventHooks.pre_call,
+        )
+        recorder: Final = RecordingLogger()
+        litellm.callbacks.append(guardrail)
+
+        response: Final = await litellm.acompletion(
+            model=CHAT_MODEL,
+            messages=CHAT_MESSAGES,
+            api_key="test-key",
+            api_base=provider.base_url,
+            callbacks=[generic_api_export.logger, recorder],
+            guardrails=[guardrail.guardrail_name],
+        )
+        await recorder.wait_for_async("async_log_success_event")
+
+        provider_messages: Final = provider.requests[0].body["messages"]
+        guard_messages: Final = guard.requests[0].body["guard_input"]["messages"]
+        exported_messages: Final = generic_api_export.exports[0].body["messages"]
+        assert response.choices[0].message.content == "Handled safely"
+        assert guard_messages == [CHAT_MESSAGES[0], CHAT_MESSAGES[3]]
+        assert provider_messages == [
+            {"role": "user", "content": [{"type": "text", "text": "Earlier safe question"}]},
+            {"role": "assistant", "content": [{"type": "text", "text": "Earlier safe answer"}]},
+            {"role": "user", "content": [{"type": "text", "text": "Employee SSN: <US_SSN>"}]},
+        ]
+        assert provider.requests[0].body["system"] == [{"type": "text", "text": "Keep the answer short"}]
+        assert exported_messages == CHAT_MESSAGES
+        assert has_rust_response_marker(response)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_rubrik_block_preserves_context_for_error_exporters(
+    backend: Backend,
+    monkeypatch: pytest.MonkeyPatch,
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    monkeypatch.setenv("RUBRIK_BATCH_SIZE", "1")
+    async with isolated_backend(backend):
+        with recording_service() as provider, recording_service() as rubrik_service:
+            provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
+            rubrik_service.expected_requests = None
+            rubrik_service.default_response = ResponseSpec(
+                body={
+                    "choices": [
+                        {
+                            "message": {
+                                "role": "assistant",
+                                "content": "Response blocked by policy",
+                                "tool_calls": [],
+                            }
+                        }
+                    ]
+                }
+            )
+            rubrik: Final = RubrikLogger(
+                api_key="test-rubrik-key",
+                api_base=rubrik_service.base_url,
+                guardrail_name="rubrik-block",
+                event_hook=GuardrailEventHooks.post_call,
+                default_on=True,
+            )
+            recorder: Final = RecordingLogger()
+            litellm.callbacks.append(rubrik)
+
+            try:
+                with pytest.raises(ModifyResponseException, match="Response blocked by policy") as raised:
+                    await litellm.acompletion(
+                        model=CHAT_MODEL,
+                        messages=CHAT_MESSAGES,
+                        api_key="test-key",
+                        api_base=provider.base_url,
+                        callbacks=[generic_api_export.logger, recorder],
+                        guardrails=[rubrik.guardrail_name],
+                    )
+                assert has_rust_response_marker(raised.value.original_response) is (backend == "rust")
+                await drain_logging()
+                await rubrik.flush_queue()
+
+                moderation_requests: Final = tuple(
+                    request for request in rubrik_service.requests if request.path == "/v1/after_completion/openai/v1"
+                )
+                batch_requests: Final = tuple(
+                    request for request in rubrik_service.requests if request.path == "/v1/litellm/batch"
+                )
+                assert len(provider.requests) == 1
+                assert len(moderation_requests) == 1
+                assert moderation_requests[0].body["request"]["messages"] == CHAT_MESSAGES
+                assert moderation_requests[0].body["response"]["choices"][0]["message"]["content"] == "Handled safely"
+                assert len(batch_requests) == 1
+                assert "Response blocked by policy" in json.dumps(batch_requests[0].body)
+                assert len(generic_api_export.exports) == 1
+                exported: Final = generic_api_export.exports[0].body
+                assert exported["status"] == "failure"
+                assert exported["error_information"]["error_class"] == "ModifyResponseException"
+                assert "Response blocked by policy" in exported["error_information"]["error_message"]
+            finally:
+                await rubrik.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_purview_audit_retains_payload_after_sdk_response(
+    backend: Backend,
+    generic_api_export: GenericAPIExportHarness,
+) -> None:
+    async with isolated_backend(backend):
+        with recording_service() as provider:
+            provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
+            purview: Final = purview_harness("purview-audit")
+            recorder: Final = RecordingLogger()
+            litellm.callbacks.append(purview.guardrail)
+            request_task: Final = asyncio.create_task(
+                litellm.acompletion(
+                    model=CHAT_MODEL,
+                    messages=CHAT_MESSAGES,
+                    api_key="test-key",
+                    api_base=provider.base_url,
+                    metadata={"user_api_key_user_id": "audit-user"},
+                    callbacks=[generic_api_export.logger, recorder],
+                    guardrails=[purview.guardrail.guardrail_name],
+                )
+            )
+            try:
+                accepted: Final = await asyncio.to_thread(purview.graph.accepted.wait, 3)
+                assert accepted
+                response: Final = await asyncio.wait_for(asyncio.shield(request_task), 2)
+                assert response.choices[0].message.content == "Handled safely"
+                assert len(await wait_for_audits(purview.graph, count=1)) == 1
+
+                purview.graph.release.set()
+                await recorder.wait_for_async("async_log_success_event")
+                audits: Final = await wait_for_audits(purview.graph)
+                activities: Final = tuple(
+                    post.body["contentToProcess"]["activityMetadata"]["activity"] for post in audits
+                )
+                assert activities == ("uploadText", "downloadText")
+                assert CHAT_MESSAGES[-1]["content"] in json.dumps(audits[0].body)
+                assert "Handled safely" in json.dumps(audits[1].body)
+                assert generic_api_export.exports[0].body["status"] == "success"
+                assert has_rust_response_marker(response) is (backend == "rust")
+            finally:
+                purview.graph.release.set()
+                if not request_task.done():
+                    request_task.cancel()
+                await asyncio.gather(request_task, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ("python", "rust"))
+async def test_purview_sync_audit_runs_without_caller_event_loop(
+    backend: Backend,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("LITERAL_BATCH_SIZE", "1")
+    async with isolated_backend(backend):
+        with recording_service() as provider, recording_service() as sink:
+            provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
+            purview: Final = purview_harness("purview-sync-audit")
+            exporter: Final = LiteralAILogger(literalai_api_key="test-key", literalai_api_url=sink.base_url)
+            litellm.callbacks.append(purview.guardrail)
+
+            try:
+                response: Final = await asyncio.to_thread(
+                    litellm.completion,
+                    model=CHAT_MODEL,
+                    messages=CHAT_MESSAGES,
+                    api_key="test-key",
+                    api_base=provider.base_url,
+                    metadata={"user_api_key_user_id": "audit-user"},
+                    callbacks=[exporter],
+                    guardrails=[purview.guardrail.guardrail_name],
+                )
+                accepted: Final = await asyncio.to_thread(purview.graph.accepted.wait, 3)
+                assert accepted
+                assert response.choices[0].message.content == "Handled safely"
+                assert len(await wait_for_audits(purview.graph, count=1)) == 1
+
+                purview.graph.release.set()
+                audits: Final = await wait_for_audits(purview.graph)
+                assert CHAT_MESSAGES[-1]["content"] in json.dumps(audits[0].body)
+                assert "Handled safely" in json.dumps(audits[1].body)
+                assert sink.requests[0].path == "/api/graphql"
+                assert "Handled safely" in json.dumps(sink.requests[0].body["variables"]["generation_0"])
+                assert has_rust_response_marker(response) is (backend == "rust")
+            finally:
+                purview.graph.release.set()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
+async def test_content_filter_post_call_blocks_provider_response(route: Route, provider: RecordingServer) -> None:
+    guardrail: Final = ContentFilterGuardrail(
+        guardrail_name="enforced-content-review",
+        event_hook=GuardrailEventHooks.post_call,
+        blocked_words=[BlockedWord(keyword=route.expected_text, action=ContentFilterAction.BLOCK)],
+    )
+    litellm.callbacks.append(guardrail)
+    with pytest.raises(HTTPException, match="Content blocked") as blocked:
+        await route.invoke(provider, guardrails=["enforced-content-review"])
+    assert blocked.value.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_azure_text_moderation_allows_messages_response_over_http(
+    recording_server: RecordingServer, otel: OtelHarness
+) -> None:
+    recording_server.expected_requests = None
+    recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_ALLOW_RESPONSE)
+    recording_server.enqueue(provider_response(MESSAGES_ROUTE))
+    guardrail: Final = azure_text_moderation(recording_server)
+    recorder: Final = RecordingLogger()
+    litellm.callbacks.append(guardrail)
+
+    response: Final = await MESSAGES_ROUTE.invoke(
+        recording_server, callbacks=[otel.logger, recorder], guardrails=[guardrail.guardrail_name]
+    )
+    await wait_for_callback(MESSAGES_ROUTE, recorder)
+
+    assert len(recording_server.requests) == 2
+    moderation_request: Final = recording_server.requests[1]
+    assert moderation_request.path == "/contentsafety/text:analyze?api-version=2024-09-01"
+    assert moderation_request.headers["ocp-apim-subscription-key"] == "test-azure-key"
+    assert moderation_request.body == {
+        "text": MESSAGES_ROUTE.expected_text,
+        "categories": ["Hate", "Sexual", "SelfHarm", "Violence"],
+        "blocklistNames": None,
+        "haltOnBlocklistHit": False,
+        "outputType": "FourSeverityLevels",
+    }
+    assert response["content"][0]["text"] == MESSAGES_ROUTE.expected_text
+    assert len(await otel.wait_for_spans()) == 1
+
+
+@pytest.mark.asyncio
+async def test_azure_text_moderation_blocks_messages_response_over_http(recording_server: RecordingServer) -> None:
+    recording_server.expected_requests = None
+    recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_BLOCK_RESPONSE)
+    recording_server.enqueue(provider_response(MESSAGES_ROUTE))
+    guardrail: Final = azure_text_moderation(recording_server)
+    litellm.callbacks.append(guardrail)
+    with pytest.raises(HTTPException, match="Violence crossed severity 2") as blocked:
+        await MESSAGES_ROUTE.invoke(recording_server, guardrails=[guardrail.guardrail_name])
+    assert blocked.value.status_code == 400
+    assert recording_server.requests[1].body["text"] == MESSAGES_ROUTE.expected_text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
+async def test_post_call_guardrail_replacement_is_what_loggers_see(
+    route: Route, provider: RecordingServer, otel: OtelHarness
+) -> None:
+    async def review(response: object) -> object:
+        match route.name:
+            case "ocr-async":
+                return response.model_copy(
+                    update={"pages": [response.pages[0].model_copy(update={"markdown": "Reviewed OCR"})]}
+                )
+            case _:
+                return {**response, "content": [{"type": "text", "text": "Reviewed Messages"}]}
+
+    guardrail: Final = ReviewGuardrail(review)
+    recorder: Final = RecordingLogger()
+    litellm.callbacks.append(guardrail)
+    response: Final = await route.invoke(provider, callbacks=[otel.logger, recorder], guardrails=["rust-review"])
+
+    assert guardrail.call_types == [CallTypes(route.call_type)]
+    event: Final = (await wait_for_callback(route, recorder))[0]
+    span: Final = (await otel.wait_for_spans())[0]
+    match route.name:
+        case "ocr-async":
+            assert response.pages[0].markdown == "Reviewed OCR"
+            assert event.response.pages[0].markdown == "Reviewed OCR"
+        case _:
+            assert response["content"][0]["text"] == "Reviewed Messages"
+            assert event.response.choices[0].message.content == "Reviewed Messages"
+            assert "Reviewed Messages" in json.dumps(dict(span.attributes))
+    assert "guardrails" not in provider.requests[0].body

tests/test_litellm_rust/messages/__init__.py

@@ -0,0 +1 @@
+

tests/test_litellm_rust/messages/test_callback_mutation.py

@@ -0,0 +1,682 @@
+import asyncio
+import copy
+import json
+import threading
+from collections.abc import Mapping
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+from tests.test_litellm_rust.support.requests import (
+    MESSAGES,
+    MESSAGES_EVENTS,
+    MESSAGES_MODEL,
+    MESSAGES_RESPONSE,
+    request_body,
+    request_headers,
+)
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def messages_server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    return recording_server
+
+
+async def call_messages(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
+    return await litellm.anthropic.messages.acreate(
+        model=MESSAGES_MODEL,
+        messages=MESSAGES,
+        max_tokens=64,
+        api_key="test-key",
+        api_base=server.base_url,
+        callbacks=callbacks,
+        **kwargs,
+    )
+
+
+@pytest.mark.asyncio
+async def test_messages_pre_call_receives_expected_provider_request(messages_server: RecordingServer) -> None:
+    observations: Final = []
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            observations.append((model, messages, copy.deepcopy(kwargs["additional_args"])))
+
+    await call_messages(messages_server, [Observe()])
+
+    assert len(observations) == 1
+    model, messages, additional_args = observations[0]
+    assert model == "claude-sonnet-4-5-20250929"
+    assert messages == MESSAGES
+    assert additional_args["api_base"] == f"{messages_server.base_url}/v1/messages"
+    assert additional_args["complete_input_dict"] == {
+        "model": "claude-sonnet-4-5-20250929",
+        "messages": MESSAGES,
+        "max_tokens": 64,
+        "stream": False,
+    }
+    assert additional_args["headers"]["x-api-key"] == "test-key"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("raise_after_edit", [False, True])
+@pytest.mark.parametrize("native", [False, True])
+async def test_messages_pre_call_edits_reach_later_callbacks_and_provider(
+    messages_server: RecordingServer, raise_after_edit: bool, native: bool, stream: bool
+) -> None:
+    litellm.rust(native)
+    observed: Final = []
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            request_body(kwargs)["temperature"] = 0.25
+            request_headers(kwargs)["x-audit-tag"] = "reviewed"
+            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)), dict(request_headers(kwargs))))
+
+    if stream:
+        messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    response: Final = await call_messages(messages_server, [Edit(), Observe()], stream=stream)
+    if stream:
+        async for _ in response:
+            pass
+
+    assert observed[0][0]["temperature"] == 0.25
+    assert observed[0][1]["x-audit-tag"] == "reviewed"
+    assert "temperature" not in messages_server.requests[0].body
+    assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed"
+
+
+@pytest.mark.asyncio
+async def test_messages_pre_call_rebinding_does_not_replace_inflight_request(
+    messages_server: RecordingServer,
+) -> 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))
+
+    await call_messages(messages_server, [Rebind(), Observe()])
+
+    assert observed == [{"replacement": True}]
+    assert messages_server.requests[0].body["messages"] == MESSAGES
+
+
+@pytest.mark.asyncio
+async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_server: RecordingServer) -> None:
+    token: Final = object()
+    observed: Final = []
+    finished: Final = asyncio.Event()
+
+    class Stash(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            kwargs["test-token"] = token
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            observed.append(kwargs["test-token"])
+            finished.set()
+
+    await call_messages(messages_server, [Stash()])
+    await asyncio.wait_for(finished.wait(), timeout=10)
+
+    assert observed == [token]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native", [False, True])
+@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"])
+@pytest.mark.parametrize("raise_after_edit", [False, True])
+async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing(
+    messages_server: RecordingServer, native: bool, provider: str, raise_after_edit: bool
+) -> None:
+    litellm.rust(native)
+    retained: Final = []
+    observed: Final = []
+
+    class Retain(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            body = request_body(kwargs)
+            message = kwargs["messages"][0]
+            assert body["messages"][0] is message
+            assert body["messages"][0]["content"] is message["content"]
+            assert body["messages"][0]["content"][0] is message["content"][0]
+            retained.append(message["content"][0])
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            retained[0]["text"] = "changed through retained reference"
+            if raise_after_edit:
+                raise RuntimeError("export failed after mutation")
+
+    class Observe(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            block = request_body(kwargs)["messages"][0]["content"][0]
+            observed.append((block is retained[0], block["text"]))
+
+    response: Final = await litellm.anthropic.messages.acreate(
+        model=MESSAGES_MODEL.replace("anthropic/", f"{provider}/"),
+        messages=[{"role": "user", "content": [{"type": "text", "text": "original"}]}],
+        max_tokens=64,
+        api_key="test-key",
+        api_base=messages_server.base_url,
+        callbacks=[Retain(), Edit(), Observe()],
+    )
+    assert observed == [(True, "changed through retained reference")]
+    assert retained[0]["text"] == "changed through retained reference"
+    assert messages_server.requests[0].body["messages"][0]["content"][0]["text"] == "original"
+    if native:
+        assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true"
+
+
+@pytest.mark.asyncio
+async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    await call_messages(messages_server, [recorder])
+    await recorder.wait_for_async("async_log_success_event")
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert recorder.names.count("async_logging_hook") == 1
+    assert recorder.names.count("async_log_success_event") == 1
+    assert "log_failure_event" not in recorder.names
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_unconsumed_stream_close_drains_and_logs_success_once(messages_server: RecordingServer) -> None:
+    messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    recorder: Final = RecordingLogger()
+    stream: Final = await call_messages(messages_server, [recorder], stream=True)
+    assert "async_log_success_event" not in recorder.names
+    await stream.aclose()
+    await stream.aclose()
+    await recorder.wait_for_async("async_log_success_event")
+    assert recorder.names.count("async_log_success_event") == 1
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_logging_drain_waits_for_suspended_callback(messages_server: RecordingServer) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    finished: Final = asyncio.Event()
+
+    class SuspendedLogger(CustomLogger):
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            started.set()
+            await release.wait()
+            finished.set()
+
+    await call_messages(messages_server, [SuspendedLogger()])
+    draining: Final = asyncio.create_task(drain_logging())
+    try:
+        await asyncio.wait_for(started.wait(), timeout=10)
+        await asyncio.sleep(0)
+        assert not draining.done()
+        assert not finished.is_set()
+    finally:
+        release.set()
+        await asyncio.wait_for(draining, timeout=10)
+
+    assert finished.is_set()
+
+
+@pytest.mark.asyncio
+async def test_messages_failure_callbacks_receive_original_provider_error(messages_server: RecordingServer) -> None:
+    messages_server.default_response = ResponseSpec(body={"error": {"message": "provider unavailable"}}, status=500)
+    messages_server.expected_requests = None
+    recorder: Final = RecordingLogger()
+
+    with pytest.raises(litellm.InternalServerError) as caught:
+        await call_messages(messages_server, [recorder])
+    events: Final = await recorder.wait_for_async("async_log_failure_event")
+
+    assert len(events) == 1
+    assert events[0].call_type == "anthropic_messages"
+    assert events[0].kwargs["exception"] is caught.value
+    assert len(messages_server.requests) == 1
+    assert "async_log_success_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_success_callbacks_receive_expected_context(messages_server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    await call_messages(
+        messages_server,
+        [recorder],
+        litellm_call_id="messages-success",
+        metadata={"source": "callback-test"},
+    )
+    async_events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert len(async_events) == 1
+    event: Final = async_events[0]
+    assert event.call_type == "anthropic_messages"
+    assert event.kwargs["litellm_call_id"] == "messages-success"
+    assert event.kwargs["litellm_params"]["metadata"]["source"] == "callback-test"
+    assert event.response.choices[0].message.content == "Hello from native Messages"
+
+
+@pytest.mark.asyncio
+async def test_messages_pre_call_runs_in_callers_execution_context(messages_server: RecordingServer) -> None:
+    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()))
+
+    caller_loop: Final = asyncio.get_running_loop()
+    await call_messages(messages_server, [Observe()])
+
+    assert observations == [(caller_loop, caller_thread)]
+
+
+@pytest.mark.asyncio
+async def test_messages_stream_logs_success_after_exhaustion(messages_server: RecordingServer) -> None:
+    messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    recorder: Final = RecordingLogger()
+    stream: Final = await call_messages(messages_server, [recorder], stream=True)
+
+    assert "async_log_success_event" not in recorder.names
+    chunks: Final = [chunk async for chunk in stream]
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert chunks
+    assert len(events) == 1
+    assert "async_log_stream_event" not in recorder.names
+    assert events[0].kwargs["complete_streaming_response"] is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native", [False, True])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_messages_compression_hook_replaces_messages_sent_to_provider(
+    messages_server: RecordingServer, native: bool, stream: bool
+) -> None:
+    litellm.rust(native)
+    compressed_messages: Final = [{"role": "user", "content": "Compressed context"}]
+    call_types: Final = []
+    recorder: Final = RecordingLogger()
+
+    class CompressMessages(CustomLogger):
+        async def async_pre_call_deployment_hook(self, kwargs, call_type):
+            assert not messages_server.requests
+            call_types.append(call_type)
+            return {**kwargs, "messages": compressed_messages}
+
+    litellm.callbacks.append(CompressMessages())
+
+    if stream:
+        messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    response: Final = await call_messages(messages_server, [recorder], stream=stream)
+    assert has_rust_response_marker(response) is native
+    if stream:
+        assert [chunk async for chunk in response]
+    await recorder.wait_for_async("async_log_success_event")
+
+    assert call_types == [CallTypes.anthropic_messages]
+    assert messages_server.requests[0].body["messages"] == compressed_messages
+    pre_calls: Final = tuple(event for event in recorder.events if event.name == "log_pre_api_call")
+    assert len(pre_calls) == 1
+    assert request_body(pre_calls[0].kwargs)["messages"] == compressed_messages
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native", [False, True])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_messages_deployment_rejection_prevents_provider_call(
+    messages_server: RecordingServer, native: bool, stream: bool
+) -> None:
+    litellm.rust(native)
+    messages_server.expected_requests = 0
+    calls: Final[list[CallTypes | None]] = []
+    recorder: Final = RecordingLogger()
+
+    class Reject(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: CallTypes | None
+        ) -> dict[str, object]:
+            calls.append(call_type)
+            raise litellm.BadRequestError(
+                message="deployment policy rejected request", model=MESSAGES_MODEL, llm_provider="anthropic"
+            )
+
+    litellm.callbacks.append(Reject())
+
+    with pytest.raises(litellm.BadRequestError, match="deployment policy rejected request"):
+        await call_messages(messages_server, [recorder], stream=stream)
+
+    assert calls == [CallTypes.anthropic_messages]
+    assert not messages_server.requests
+    assert "log_pre_api_call" not in recorder.names
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native", [False, True])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_messages_provider_open_failure_notifies_deployment_once_and_preserves_error(
+    messages_server: RecordingServer, native: bool, stream: bool
+) -> None:
+    litellm.rust(native)
+    messages_server.enqueue(ResponseSpec(body={"error": {"message": "provider unavailable"}}, status=500))
+    calls: Final[list[tuple[CallTypes | None, Exception, int]]] = []
+
+    class FailingObserver(CustomLogger):
+        async def async_post_call_failure_deployment_hook(
+            self,
+            request_data: Mapping[str, object],
+            exception: Exception,
+            call_type: CallTypes | None,
+            fallback_depth: int | None = None,
+        ) -> None:
+            calls.append((call_type, exception, len(messages_server.requests)))
+            raise RuntimeError("deployment observer unavailable")
+
+    litellm.callbacks.append(FailingObserver())
+    recorder: Final = RecordingLogger()
+
+    with pytest.raises(litellm.InternalServerError) as raised:
+        await call_messages(messages_server, [recorder], stream=stream)
+    await recorder.wait_for_async("async_log_failure_event")
+
+    assert len(calls) == 1
+    call_type, exception, request_count = calls[0]
+    assert call_type == CallTypes.anthropic_messages
+    assert isinstance(exception, litellm.InternalServerError)
+    assert exception.status_code == raised.value.status_code == 500
+    assert request_count == len(messages_server.requests) == 1
+    assert recorder.names.count("async_log_failure_event") == 1
+    assert "async_log_success_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_post_call_guardrail_replacement_reaches_caller_and_logging(
+    messages_server: RecordingServer,
+) -> None:
+    recorder: Final = RecordingLogger()
+    call_types: Final = []
+
+    class ReviewResponse(CustomLogger):
+        async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+            call_types.append(call_type)
+            response["content"][0]["text"] = "Reviewed response"
+            return response
+
+    litellm.callbacks.append(ReviewResponse())
+
+    response: Final = await call_messages(messages_server, [recorder])
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert call_types == [CallTypes.anthropic_messages]
+    assert response["content"][0]["text"] == "Reviewed response"
+    assert events[0].response.choices[0].message.content == "Reviewed response"
+
+
+@pytest.mark.asyncio
+async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None:
+    observations: Final = []
+    exported: Final = asyncio.Event()
+
+    class RecordGuardrailVerdict(CustomLogger):
+        async def async_logging_hook(self, kwargs, result, call_type):
+            observations.append("guardrail")
+            return {**kwargs, "guardrail-verdict": "allowed"}, result
+
+    class ExportLog(CustomLogger):
+        async def async_logging_hook(self, kwargs, result, call_type):
+            return kwargs, result
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            observations.append(("export", kwargs["guardrail-verdict"]))
+            exported.set()
+
+    response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()])
+    await asyncio.wait_for(exported.wait(), timeout=10)
+
+    assert observations == ["guardrail", ("export", "allowed")]
+    assert response["content"][0]["text"] == "Hello from native Messages"
+
+
+@pytest.mark.asyncio
+async def test_messages_success_callback_failure_does_not_skip_later_loggers(
+    messages_server: RecordingServer,
+) -> None:
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    response: Final = await call_messages(messages_server, [UnavailableExporter(), recorder])
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert response["content"][0]["text"] == "Hello from native Messages"
+    assert len(events) == 1
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_stream_success_callback_failure_does_not_skip_later_loggers(
+    messages_server: RecordingServer,
+) -> None:
+    messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    stream: Final = await call_messages(messages_server, [UnavailableExporter(), recorder], stream=True)
+    chunks: Final = [chunk async for chunk in stream]
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert chunks
+    assert len(events) == 1
+    assert events[0].response.choices[0].message.content == "Hello from native Messages"
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_stream_failure_callback_failure_does_not_skip_later_loggers(
+    messages_server: RecordingServer,
+) -> None:
+    messages_server.enqueue(
+        ResponseSpec(
+            body=None,
+            events=(
+                MESSAGES_EVENTS[0],
+                (
+                    "error",
+                    {"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}},
+                ),
+            ),
+        )
+    )
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    stream: Final = await call_messages(messages_server, [UnavailableExporter(), recorder], stream=True)
+    with pytest.raises(litellm.APIError) as caught:
+        async for _ in stream:
+            pass
+    events: Final = await recorder.wait_for_async("async_log_failure_event")
+
+    assert len(events) == 1
+    assert isinstance(events[0].kwargs["exception"], litellm.APIError)
+    assert type(events[0].kwargs["exception"]) is type(caught.value)
+    assert "async_log_success_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_messages_concurrent_calls_keep_callback_state_isolated(messages_server: RecordingServer) -> None:
+    messages_server.expected_requests = 4
+    tokens: Final = {f"messages-{index}": object() for index in range(4)}
+    terminal_state: Final = []
+
+    class CorrelateCallState(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]]
+            super().log_pre_api_call(model, messages, kwargs)
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"]))
+            await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
+
+    correlate: Final = CorrelateCallState()
+
+    await asyncio.gather(*(call_messages(messages_server, [correlate], litellm_call_id=call_id) for call_id in tokens))
+    await correlate.wait_for_async("async_log_success_event", count=4)
+
+    assert len(terminal_state) == 4
+    assert all(token is tokens[call_id] for call_id, token in terminal_state)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_messages_cancelled_call_runs_no_terminal_callbacks(
+    messages_server: RecordingServer, stream: bool
+) -> None:
+    messages_server.default_response = ResponseSpec(
+        body=MESSAGES_RESPONSE,
+        delay=0.5,
+        events=MESSAGES_EVENTS if stream else (),
+    )
+    recorder: Final = RecordingLogger()
+    task: Final = asyncio.create_task(call_messages(messages_server, [recorder], stream=stream))
+
+    async with asyncio.timeout(10):
+        while not messages_server.requests:
+            await asyncio.sleep(0.01)
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert "log_success_event" not in recorder.names
+    assert "async_log_success_event" not in recorder.names
+    assert "log_failure_event" not in recorder.names
+    assert "async_log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status", [200, 500])
+async def test_whole_call_with_supplied_logger_still_owns_terminal_callbacks(
+    messages_server: RecordingServer, status: int
+) -> None:
+    from litellm.rust_bridge import messages as bridge
+
+    events: Final = []
+
+    class SuppliedLogger:
+        def pre_call(self, **kwargs):
+            events.append(("pre", kwargs))
+
+        async def async_success_handler(self, response, start, end):
+            events.append(("success", response))
+
+        def _should_run_sync_callbacks_for_async_calls(self):
+            return False
+
+        def failure_handler(self, error, trace, start, end):
+            events.append(("sync_failure", error))
+
+        async def async_failure_handler(self, error, trace, start, end):
+            events.append(("async_failure", error))
+
+    messages_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, status=status)
+    arguments: Final = {
+        "model": "claude-opus-5",
+        "body": {"model": "claude-opus-5", "messages": MESSAGES, "max_tokens": 64},
+        "api_key": "test",
+        "api_base": messages_server.base_url,
+        "custom_llm_provider": "anthropic",
+        "extra_headers": {},
+        "timeout": 5.0,
+        "logging_obj": SuppliedLogger(),
+    }
+    if status == 200:
+        response: Final = await bridge.amessages(**arguments)
+        await drain_logging()
+        assert [name for name, value in events] == ["pre", "success"]
+        assert events[1][1] is response
+        return
+
+    with pytest.raises(litellm.InternalServerError) as raised:
+        await bridge.amessages(**arguments)
+    assert [name for name, value in events] == ["pre", "sync_failure", "async_failure"]
+    assert events[1][1] is raised.value
+    assert events[2][1] is raised.value
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("detach", (False, True))
+async def test_direct_bridge_stream_owns_callbacks(messages_server: RecordingServer, detach: bool) -> None:
+    from litellm.rust_bridge import messages as bridge
+
+    recorder: Final = RecordingLogger()
+    release: Final = threading.Event()
+    if not detach:
+        release.set()
+    messages_server.enqueue(
+        ResponseSpec(
+            body=None,
+            chunks=tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS),
+            release=release,
+        )
+    )
+    stream: Final = await bridge.amessages(
+        model=MESSAGES_MODEL.split("/", 1)[-1],
+        body={"model": MESSAGES_MODEL.split("/", 1)[-1], "messages": MESSAGES, "max_tokens": 64, "stream": True},
+        api_key="test-key",
+        api_base=messages_server.base_url,
+        custom_llm_provider="anthropic",
+        extra_headers=None,
+        timeout=5.0,
+        request_arguments={"callbacks": [recorder]},
+    )
+    assert "async_log_success_event" not in recorder.names
+    try:
+        if detach:
+            assert await anext(stream)
+            await asyncio.wait_for(stream.aclose(), timeout=1)
+            await stream.aclose()
+            assert "async_log_success_event" not in recorder.names
+            assert "async_log_failure_event" not in recorder.names
+        else:
+            assert b"message_stop" in b"".join([chunk async for chunk in stream])
+    finally:
+        release.set()
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+    assert len(events) == 1
+    assert events[0].stream is True
+    assert events[0].response.usage.total_tokens == 9
+    assert events[0].response.choices[0].message.content == "Hello from native Messages"
+    assert recorder.names.count("log_pre_api_call") == 1
+    await stream.aclose()

tests/test_litellm_rust/messages/test_streaming.py

@@ -0,0 +1,275 @@
+import json
+from collections.abc import AsyncIterator
+from typing import Final, cast
+
+import pytest
+
+import litellm
+from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def messages_server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    return recording_server
+
+
+async def call_messages(server: RecordingServer, **kwargs: object):
+    return await litellm.anthropic.messages.acreate(
+        model=MESSAGES_MODEL,
+        messages=MESSAGES,
+        max_tokens=64,
+        api_key="test-key",
+        api_base=server.base_url,
+        **kwargs,
+    )
+
+
+def assert_native_request(server: RecordingServer) -> None:
+    assert len(server.requests) == 1
+    assert "accept-encoding" not in server.requests[0].headers
+
+
+def assert_native_response(response: object) -> None:
+    assert isinstance(response, dict)
+    assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
+
+
+@pytest.mark.asyncio
+async def test_messages_sends_expected_provider_request(messages_server: RecordingServer) -> None:
+    response: Final = await call_messages(messages_server)
+
+    assert response["content"] == [{"type": "text", "text": "Hello from native Messages"}]
+    assert_native_response(response)
+    assert_native_request(messages_server)
+    request: Final = messages_server.requests[0]
+    assert request.path == "/v1/messages"
+    assert request.headers["x-api-key"] == "test-key"
+    assert request.headers["anthropic-version"] == "2023-06-01"
+    assert request.body == {
+        "model": "claude-sonnet-4-5-20250929",
+        "messages": MESSAGES,
+        "max_tokens": 64,
+    }
+
+
+@pytest.mark.asyncio
+async def test_messages_sends_custom_headers(messages_server: RecordingServer) -> None:
+    response: Final = await call_messages(messages_server, extra_headers={"x-trace-id": "trace-1"})
+
+    assert_native_response(response)
+    assert messages_server.requests[0].headers["x-trace-id"] == "trace-1"
+
+
+@pytest.mark.asyncio
+async def test_messages_resolves_provider_credentials(
+    messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key")
+
+    response: Final = await litellm.anthropic.messages.acreate(
+        model=MESSAGES_MODEL,
+        messages=MESSAGES,
+        max_tokens=64,
+        api_base=messages_server.base_url,
+    )
+
+    assert_native_response(response)
+    assert messages_server.requests[0].headers["x-api-key"] == "environment-key"
+
+
+@pytest.mark.asyncio
+async def test_messages_explicit_credentials_override_defaults(
+    messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key")
+
+    response: Final = await call_messages(messages_server)
+
+    assert_native_response(response)
+    assert messages_server.requests[0].headers["x-api-key"] == "test-key"
+
+
+@pytest.mark.asyncio
+async def test_azure_messages_uses_foundry_endpoint_and_credentials(messages_server: RecordingServer) -> None:
+    response: Final = await litellm.anthropic.messages.acreate(
+        model="azure_ai/claude-opus-4.5",
+        messages=MESSAGES,
+        max_tokens=64,
+        api_key="azure-key",
+        api_base=f"{messages_server.base_url}/anthropic",
+    )
+
+    assert_native_response(response)
+    assert_native_request(messages_server)
+    assert messages_server.requests[0].path == "/anthropic/v1/messages"
+    assert messages_server.requests[0].headers["x-api-key"] == "azure-key"
+
+
+@pytest.mark.asyncio
+async def test_messages_stream_yields_anthropic_events(messages_server: RecordingServer) -> None:
+    messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS))
+    stream: Final = cast(AsyncIterator[bytes], await call_messages(messages_server, stream=True))
+    payload: Final = b"".join([chunk async for chunk in stream])
+
+    assert_native_request(messages_server)
+    assert messages_server.requests[0].body["stream"] is True
+    assert b"event: message_start" in payload
+    assert b"event: content_block_delta" in payload
+    assert b"Hello from native Messages" in payload
+    assert b"event: message_stop" in payload
+    message_delta: Final = next(
+        json.loads(block.split(b"data: ", 1)[1])
+        for block in payload.split(b"\n\n")
+        if block.startswith(b"event: message_delta")
+    )
+    assert message_delta["usage"] == {"input_tokens": 5, "output_tokens": 4}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"])
+async def test_messages_delivers_before_upstream_finishes(messages_server: RecordingServer, provider: str) -> None:
+    import asyncio
+    import threading
+
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+
+    release: Final = threading.Event()
+    recorder: Final = RecordingLogger()
+    chunks: Final = tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS)
+    messages_server.enqueue(ResponseSpec(body=None, chunks=chunks, release=release))
+    try:
+        stream: Final = await asyncio.wait_for(
+            litellm.anthropic.messages.acreate(
+                model=MESSAGES_MODEL.split("/", 1)[-1],
+                custom_llm_provider=provider,
+                messages=MESSAGES,
+                max_tokens=64,
+                api_key="test-key",
+                api_base=messages_server.base_url,
+                stream=True,
+                callbacks=[recorder],
+            ),
+            timeout=5,
+        )
+        first: Final = await asyncio.wait_for(anext(stream), timeout=2)
+        assert b"message_start" in first
+        assert not release.is_set()
+        assert "async_log_success_event" not in recorder.names
+        assert messages_server.requests[0].body["stream"] is True
+        assert stream._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
+        assert messages_server.requests[0].path == (
+            "/v1/messages" if provider == "anthropic" else "/anthropic/v1/messages"
+        )
+        release.set()
+        remaining: Final = b"".join([chunk async for chunk in stream])
+        assert first + remaining == b"".join(chunks)
+        await stream.aclose()
+        events: Final = await recorder.wait_for_async("async_log_success_event")
+        assert len(events) == 1
+        assert "async_log_failure_event" not in recorder.names
+        assert events[0].response.choices[0].message.content == "Hello from native Messages"
+        assert events[0].response.usage.completion_tokens == 4
+        assert len(messages_server.requests) == 1
+    finally:
+        release.set()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ending", ["truncated", "provider_error", "http_error"])
+async def test_messages_stream_failures_do_not_replay(messages_server: RecordingServer, ending: str) -> None:
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+
+    recorder: Final = RecordingLogger()
+    events: Final = (
+        MESSAGES_EVENTS[:-1]
+        if ending == "truncated"
+        else (
+            *MESSAGES_EVENTS[:1],
+            ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}),
+        )
+    )
+    messages_server.enqueue(
+        ResponseSpec(
+            body={"error": "rejected"},
+            status=429 if ending == "http_error" else 200,
+            events=() if ending == "http_error" else events,
+        )
+    )
+
+    async def consume() -> None:
+        stream: Final = await call_messages(messages_server, stream=True, callbacks=[recorder])
+        async for _ in stream:
+            pass
+
+    with pytest.raises(litellm.APIError):
+        await consume()
+    await recorder.wait_for_async("async_log_failure_event")
+    assert recorder.names.count("async_log_failure_event") == 1
+    assert "async_log_success_event" not in recorder.names
+    assert len(messages_server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("abandon", ["close", "cancel", "drop", "exhaust"])
+async def test_messages_stream_abandonment_releases_roots(messages_server: RecordingServer, abandon: str) -> None:
+    import asyncio
+    import gc
+    import threading
+    import weakref
+
+    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+
+    class NamesOnlyLogger(RecordingLogger):
+        def _record(self, name: str, kwargs: object = None, response: object = None) -> None:
+            super()._record(name)
+
+    class Root:
+        pass
+
+    root = Root()
+    reference: Final = weakref.ref(root)
+    release: Final = threading.Event()
+    recorder: Final = NamesOnlyLogger()
+    chunks: Final = tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS)
+    messages_server.enqueue(ResponseSpec(body=None, chunks=chunks, release=release))
+    try:
+        stream = await call_messages(messages_server, stream=True, callbacks=[recorder], metadata={"retained": root})
+        del root
+        assert reference() is not None
+        assert b"message_start" in await anext(stream)
+        if abandon == "exhaust":
+            release.set()
+            async for _ in stream:
+                pass
+            await recorder.wait_for_async("async_log_success_event")
+            gc.collect()
+            assert reference() is None
+            assert recorder.names.count("async_log_success_event") == 1
+            return
+        if abandon == "cancel":
+            task: Final = asyncio.create_task(anext(stream))
+            await asyncio.sleep(0.05)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        elif abandon == "close":
+            await stream.aclose()
+            await stream.aclose()
+        del stream
+        gc.collect()
+        assert reference() is not None
+        assert "async_log_success_event" not in recorder.names
+        assert "async_log_failure_event" not in recorder.names
+        release.set()
+        await recorder.wait_for_async("async_log_success_event")
+        assert recorder.names.count("async_log_success_event") == 1
+        assert "async_log_failure_event" not in recorder.names
+        assert len(messages_server.requests) == 1
+        gc.collect()
+        assert reference() is None
+    finally:
+        release.set()

tests/test_litellm_rust/ocr/__init__.py

@@ -0,0 +1 @@
+

tests/test_litellm_rust/ocr/test_callbacks.py

@@ -0,0 +1,502 @@
+import asyncio
+import copy
+import queue
+import threading
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+from tests.test_litellm_rust.support.requests import (
+    OCR_DOCUMENT,
+    OCR_RESPONSE,
+    call_native_aocr,
+    call_native_ocr,
+    request_body,
+    request_headers,
+)
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+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
+
+
+def call_ocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
+    return call_native_ocr(server, callbacks=callbacks, **kwargs)
+
+
+async def call_aocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
+    return await call_native_aocr(server, callbacks=callbacks, **kwargs)
+
+
+def test_pre_call_receives_expected_provider_request(ocr_server: RecordingServer) -> 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": OCR_DOCUMENT,
+        "pages": [0],
+    }
+
+
+@pytest.mark.parametrize("raise_after_edit", [False, True])
+def test_pre_call_body_edits_reach_later_callbacks_and_provider(
+    ocr_server: RecordingServer, 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: RecordingServer) -> 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"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rust_enabled", [False, True])
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_pre_call_nested_mutation_updates_retained_references(
+    ocr_server: RecordingServer, rust_enabled: bool, asynchronous: bool
+) -> None:
+    litellm.rust(rust_enabled)
+    original: Final = dict(OCR_DOCUMENT)
+    replacement_url: Final = "data:application/pdf;base64,ZGVm"
+    retained: Final = []
+
+    class Retain(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            assert request_body(kwargs)["document"] is original
+            retained.append(request_body(kwargs)["document"])
+
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            original["document_url"] = replacement_url
+
+    arguments: Final = {
+        "model": "mistral/mistral-ocr-latest",
+        "document": original,
+        "api_key": "test-key",
+        "api_base": ocr_server.base_url,
+        "callbacks": [Retain(), Edit()],
+    }
+    response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
+
+    assert retained[0]["document_url"] == replacement_url
+    assert original["document_url"] == replacement_url
+    assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
+    assert has_rust_response_marker(response) is rust_enabled
+
+
+def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None:
+    original: Final = dict(OCR_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
+
+    call_native_ocr(
+        ocr_server,
+        document=original,
+        callbacks=[RetainAndReplace()],
+    )
+
+    assert retained[0] is original
+    assert original["document_url"] == OCR_DOCUMENT["document_url"]
+    assert ocr_server.requests[0].body["document"] == replacement
+
+
+def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: RecordingServer) -> 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": OCR_DOCUMENT}
+
+
+def test_queued_payload_observes_later_callback_mutations(ocr_server: RecordingServer) -> 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_pre_call_state_reaches_terminal_callbacks(ocr_server: RecordingServer) -> 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: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    await call_aocr(
+        ocr_server,
+        [recorder],
+        litellm_call_id="ocr-success",
+        metadata={"source": "callback-test"},
+    )
+    events: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert len(events) == 1
+    assert events[0].call_type == "aocr"
+    assert events[0].kwargs["litellm_call_id"] == "ocr-success"
+    assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test"
+    assert events[0].response.pages[0].markdown == "native OCR response"
+
+
+@pytest.mark.asyncio
+async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: RecordingServer) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    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)
+
+
+@pytest.mark.asyncio
+async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingServer) -> None:
+    caller_loop: Final = asyncio.get_running_loop()
+    caller_thread: Final = threading.current_thread()
+    recorder: Final = RecordingLogger()
+
+    await call_aocr(ocr_server, [recorder])
+
+    events: Final = await recorder.wait_for_async("log_pre_api_call")
+    assert len(events) == 1
+    assert events[0].loop is caller_loop
+    assert events[0].thread is caller_thread
+
+
+@pytest.mark.asyncio
+async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: RecordingServer) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    token: Final = object()
+    observed: Final = []
+
+    class TrackInFlightRequest(CustomLogger):
+        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):
+            observed.append(("sync", kwargs["request-token"]))
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            observed.append(("async", kwargs["request-token"]))
+
+    with pytest.raises(litellm.InternalServerError):
+        await call_aocr(ocr_server, [TrackInFlightRequest()])
+
+    assert [event for event, _ in observed] == ["sync", "async"]
+    assert all(observed_token is token for _, observed_token in observed)
+
+
+@pytest.mark.asyncio
+async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_callbacks(
+    ocr_server: RecordingServer,
+) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    recorder: Final = RecordingLogger()
+
+    class UnavailableExporter(CustomLogger):
+        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            raise RuntimeError("exporter unavailable")
+
+    with pytest.raises(litellm.InternalServerError) as caught:
+        await call_aocr(ocr_server, [UnavailableExporter(), recorder])
+
+    sync_events: Final = tuple(event for event in recorder.events if event.name == "log_failure_event")
+    async_events: Final = tuple(event for event in recorder.events if event.name == "async_log_failure_event")
+    assert len(sync_events) == 1
+    assert len(async_events) == 1
+    assert sync_events[0].kwargs["exception"] is caught.value
+    assert async_events[0].kwargs["exception"] is caught.value
+    assert "async_log_success_event" not in recorder.names
+
+
+def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    call_ocr(
+        ocr_server,
+        [recorder, recorder],
+        success_callback=[recorder],
+        failure_callback=[recorder],
+    )
+    recorder.wait_for("log_success_event")
+
+    assert recorder.names.count("log_pre_api_call") == 1
+    assert recorder.names.count("logging_hook") == 1
+    assert recorder.names.count("log_success_event") == 1
+    assert "log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_azure_token_callback_precedes_logger_and_preserves_context(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+) -> None:
+    from contextvars import ContextVar
+    from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr
+
+    context: Final = ContextVar("azure-token-context", default="missing")
+    context.set("caller")
+    caller_thread: Final = threading.current_thread()
+    caller_loop: Final = asyncio.get_running_loop()
+    observations: Final = []
+
+    class Provider:
+        def __call__(self) -> str:
+            assert context.get() == "caller"
+            assert threading.current_thread() is caller_thread
+            assert asyncio.get_running_loop() is caller_loop
+            observations.append("token")
+            return "caller-token"
+
+    class Edit(CustomLogger):
+        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"
+
+    provider: Final = Provider()
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": provider,
+        "callbacks": [Edit()],
+    }
+    response: Final = (
+        await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments)
+    )
+    assert has_rust_response_marker(response)
+    assert observations == ["token", "pre_call"]
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer edited"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_azure_token_callback_can_reenter_native_sdk(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+) -> None:
+    from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr
+
+    ocr_server.expected_requests = 2
+    calls: Final = []
+
+    def provider() -> str:
+        calls.append("token")
+        nested: Final = public_ocr(ocr_server)
+        assert has_rust_response_marker(nested)
+        return "outer-token"
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": provider,
+    }
+    response: Final = (
+        await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments)
+    )
+    assert has_rust_response_marker(response)
+    assert calls == ["token"]
+    assert [request.headers["authorization"] for request in ocr_server.requests] == [
+        "Bearer test-key",
+        "Bearer outer-token",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_concurrent_azure_token_callbacks_keep_results_and_errors_separate(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+) -> None:
+    from tests.test_litellm_rust.support.requests import call_aocr as public_aocr
+
+    ocr_server.expected_requests = 2
+
+    async def request(token: str, fail: bool) -> object:
+        def provider() -> str:
+            if fail:
+                raise ValueError(token)
+            return token
+
+        return await public_aocr(
+            ocr_server,
+            model="azure_ai/mistral-ocr-latest",
+            api_key=None,
+            azure_ad_token_provider=provider,
+        )
+
+    responses: Final = await asyncio.gather(
+        request("first", False),
+        request("failed", True),
+        request("second", False),
+        return_exceptions=True,
+    )
+    assert has_rust_response_marker(responses[0])
+    assert isinstance(responses[1], litellm.APIConnectionError)
+    assert "Failed to get Azure AD token: failed" in str(responses[1])
+    assert has_rust_response_marker(responses[2])
+    assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [
+        "Bearer first",
+        "Bearer second",
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"])
+async def test_azure_token_provider_is_released_after_request(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    outcome: str,
+) -> None:
+    import gc
+    import weakref
+    from tests.test_litellm_rust.support.callback_recorder import drain_logging
+    from tests.test_litellm_rust.support.requests import call_aocr as public_aocr
+
+    class Provider:
+        def __call__(self) -> str:
+            if outcome == "failure":
+                raise ValueError("unavailable")
+            return "caller-token"
+
+    async def invoke() -> weakref.ReferenceType[Provider]:
+        provider: Final = Provider()
+        reference: Final = weakref.ref(provider)
+        if outcome == "failure":
+            ocr_server.expected_requests = 0
+            with pytest.raises(litellm.APIConnectionError):
+                await public_aocr(
+                    ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider
+                )
+        elif outcome == "cancellation":
+            ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
+            task: Final = asyncio.create_task(
+                public_aocr(
+                    ocr_server,
+                    model="azure_ai/mistral-ocr-latest",
+                    api_key=None,
+                    azure_ad_token_provider=provider,
+                )
+            )
+            await ocr_server.wait_for_requests(1)
+            assert reference() is provider
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        else:
+            response: Final = await public_aocr(
+                ocr_server,
+                model="azure_ai/mistral-ocr-latest",
+                api_key=None,
+                azure_ad_token_provider=provider,
+            )
+            assert has_rust_response_marker(response)
+        return reference
+
+    reference: Final = await invoke()
+    await drain_logging()
+    await asyncio.sleep(0)
+    gc.collect()
+    assert reference() is None

tests/test_litellm_rust/ocr/test_dispatch.py

@@ -0,0 +1,78 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+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
+
+
+def test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None:
+    response: Final = litellm.ocr(
+        model=OCR_MODEL,
+        document=OCR_DOCUMENT,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert has_rust_response_marker(response)
+
+
+@pytest.mark.asyncio
+async def test_public_aocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None:
+    response: Final = await litellm.aocr(
+        model=OCR_MODEL,
+        document=OCR_DOCUMENT,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert has_rust_response_marker(response)
+
+
+def test_public_ocr_falls_back_when_native_transport_declines(ocr_server: RecordingServer) -> None:
+    response: Final = litellm.ocr(
+        model=OCR_MODEL,
+        document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert not has_rust_response_marker(response)
+
+
+def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingServer) -> None:
+    litellm.rust(False)
+
+    response: Final = litellm.ocr(
+        model=OCR_MODEL,
+        document=OCR_DOCUMENT,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert not has_rust_response_marker(response)
+
+
+def test_native_ocr_requests_an_uncompressed_response(ocr_server: RecordingServer) -> None:
+    litellm.ocr(
+        model=OCR_MODEL,
+        document=OCR_DOCUMENT,
+        api_key="test-key",
+        api_base=ocr_server.base_url,
+    )
+
+    assert ocr_server.requests[0].headers["accept-encoding"] == "identity"

tests/test_litellm_rust/ocr/test_requests.py

@@ -0,0 +1,447 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.support.requests import (
+    OCR_DOCUMENT,
+    OCR_MODEL,
+    OCR_RESPONSE,
+    call_native_ocr,
+    call_aocr,
+    call_ocr as call_public_ocr,
+)
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.provenance import has_rust_response_marker
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_azure_ocr_calls_python_token_provider(
+    ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
+) -> None:
+    calls: Final = []
+
+    def token_provider() -> str:
+        calls.append("token")
+        return "callback-token"
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": token_provider,
+    }
+    response: Final = (
+        await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
+    )
+
+    assert calls == ["token"]
+    assert has_rust_response_marker(response)
+    assert response.pages[0].markdown == "native OCR response"
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token"
+
+
+@pytest.fixture
+def ocr_server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
+    return recording_server
+
+
+def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    return call_native_ocr(server, **kwargs)
+
+
+def assert_native_request(server: RecordingServer) -> 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: RecordingServer) -> None:
+    response: Final = call_ocr(ocr_server)
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].path == "/v1/ocr"
+    assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
+
+
+def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: RecordingServer) -> None:
+    ocr_server.expected_requests = 0
+    recorder: Final = RecordingLogger()
+
+    with pytest.raises(NotImplementedError, match="OCR file document preparation"):
+        call_native_ocr(
+            ocr_server,
+            document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
+            callbacks=[recorder],
+        )
+
+    assert ocr_server.requests == []
+    assert recorder.events == ()
+
+
+def test_ocr_sends_optional_parameters(ocr_server: RecordingServer) -> 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: RecordingServer) -> 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: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
+
+    call_native_ocr(ocr_server, api_key=None)
+
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key"
+
+
+def test_ocr_explicit_credentials_override_defaults(
+    ocr_server: RecordingServer, 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: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key")
+    monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url)
+
+    call_native_ocr(ocr_server, model="azure_ai/pixtral-12b-2409", api_key=None, api_base=None)
+
+    assert_native_request(ocr_server)
+    assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr"
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer azure-key"
+
+
+def test_ocr_resolves_vertex_project_and_location(ocr_server: RecordingServer) -> None:
+    call_native_ocr(
+        ocr_server,
+        model="vertex_ai/mistral-ocr-2505",
+        api_key="vertex-token",
+        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: RecordingServer) -> 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: RecordingServer) -> None:
+    ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
+
+    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: RecordingServer) -> None:
+    ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
+
+    with pytest.raises(RuntimeError, match="OCR transport failed"):
+        call_ocr(ocr_server, timeout=0.01)
+
+    assert len(ocr_server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("backend", ["python", "rust"])
+@pytest.mark.parametrize(
+    "credentials, expected_token, expected_calls",
+    [
+        ({"api_key": "resource-key"}, "resource-key", 0),
+        ({"azure_ad_token": "static-token"}, "callback-1", 1),
+        ({"extra_headers": {"Authorization": "Bearer override"}}, "override", 1),
+    ],
+)
+async def test_azure_ocr_token_provider_precedence(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+    backend: str,
+    credentials: dict[str, object],
+    expected_token: str,
+    expected_calls: int,
+) -> None:
+    litellm.rust(backend == "rust")
+    calls: Final = []
+
+    def token_provider() -> str:
+        calls.append("token")
+        return f"callback-{len(calls)}"
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": token_provider,
+        **credentials,
+    }
+    response: Final = (
+        await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
+    )
+    assert has_rust_response_marker(response) == (backend == "rust")
+    assert len(calls) == expected_calls
+    assert len(ocr_server.requests) == 1
+    assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_token}"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_azure_ocr_does_not_cache_caller_tokens(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+) -> None:
+    calls: Final = []
+    ocr_server.expected_requests = 2
+
+    def token_provider() -> str:
+        calls.append("token")
+        return f"callback-{len(calls)}"
+
+    for _ in range(2):
+        arguments: Final = {
+            "model": "azure_ai/mistral-ocr-latest",
+            "api_key": None,
+            "azure_ad_token_provider": token_provider,
+        }
+        response: Final = (
+            await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
+        )
+        assert has_rust_response_marker(response)
+    assert len(calls) == 2
+    assert [request.headers["authorization"] for request in ocr_server.requests] == [
+        "Bearer callback-1",
+        "Bearer callback-2",
+    ]
+
+
+class TokenAbort(BaseException):
+    pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("backend", ["python", "rust"])
+@pytest.mark.parametrize("failure", ["non_string", "type_error", "ordinary", "abort"])
+async def test_azure_ocr_token_failure_stops_execution(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+    backend: str,
+    failure: str,
+) -> None:
+    litellm.rust(backend == "rust")
+    ocr_server.expected_requests = 0
+    calls: Final = []
+    recorder: Final = RecordingLogger()
+    original: Final = {
+        "type_error": TypeError("token type"),
+        "ordinary": ValueError("token unavailable"),
+        "abort": TokenAbort("abort"),
+    }
+
+    def token_provider() -> object:
+        calls.append("token")
+        if failure == "non_string":
+            return 123
+        raise original[failure]
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": token_provider,
+        "callbacks": [recorder],
+    }
+    expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
+    with pytest.raises(expected) as caught:
+        await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
+    assert calls == ["token"]
+    assert ocr_server.requests == []
+    assert "log_pre_api_call" not in recorder.names
+    if failure == "ordinary":
+        assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
+        assert isinstance(caught.value.__context__, RuntimeError)
+        assert caught.value.__context__.__cause__ is original[failure]
+    elif failure == "abort":
+        assert caught.value is original[failure]
+    elif failure == "type_error":
+        assert caught.value.__context__ is original[failure]
+    else:
+        assert isinstance(caught.value.__context__, TypeError)
+
+
+@pytest.mark.parametrize(
+    "configuration",
+    [
+        {"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"},
+        {"model": "azure_ai/doc-intelligence/prebuilt-read"},
+        {"document": {"type": "file", "file": b"pdf"}},
+    ],
+)
+def test_azure_unsupported_native_auth_does_not_invoke_provider(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    configuration: dict[str, object],
+) -> None:
+    ocr_server.expected_requests = 0
+    calls: Final = []
+    recorder: Final = RecordingLogger()
+
+    def provider() -> str:
+        calls.append("token")
+        return "unused"
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": provider,
+        "callbacks": [recorder],
+        **configuration,
+    }
+    with pytest.raises(NotImplementedError):
+        call_native_ocr(ocr_server, **arguments)
+    assert calls == []
+    assert recorder.events == ()
+    assert ocr_server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["python", "rust"])
+async def test_azure_missing_endpoint_prevents_token_callback(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    backend: str,
+) -> None:
+    litellm.rust(backend == "rust")
+    ocr_server.expected_requests = 0
+    calls: Final = []
+
+    def provider() -> str:
+        calls.append("token")
+        return "unused"
+
+    with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI API Base"):
+        await call_aocr(
+            ocr_server,
+            model="azure_ai/mistral-ocr-latest",
+            api_key=None,
+            api_base=None,
+            azure_ad_token_provider=provider,
+        )
+    assert calls == []
+    assert ocr_server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["python", "rust"])
+async def test_azure_empty_callback_token_does_not_restore_static_token(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    backend: str,
+) -> None:
+    litellm.rust(backend == "rust")
+    ocr_server.expected_requests = 0
+
+    def provider() -> str:
+        return ""
+
+    with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"):
+        await call_aocr(
+            ocr_server,
+            model="azure_ai/mistral-ocr-latest",
+            api_key=None,
+            azure_ad_token="static-token",
+            azure_ad_token_provider=provider,
+        )
+    assert ocr_server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["python", "rust"])
+async def test_azure_falsey_callable_leaves_static_token_unchanged(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    backend: str,
+) -> None:
+    litellm.rust(backend == "rust")
+    calls: Final = []
+
+    class Provider:
+        def __bool__(self) -> bool:
+            return False
+
+        def __call__(self) -> str:
+            calls.append("token")
+            return "unused"
+
+    response: Final = await call_aocr(
+        ocr_server,
+        model="azure_ai/mistral-ocr-latest",
+        api_key=None,
+        azure_ad_token="static-token",
+        azure_ad_token_provider=Provider(),
+    )
+    assert has_rust_response_marker(response) == (backend == "rust")
+    assert calls == []
+    assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["python", "rust"])
+async def test_azure_token_callback_does_not_await_coroutine_result(
+    ocr_server: RecordingServer,
+    isolated_azure_auth: None,
+    backend: str,
+) -> None:
+    litellm.rust(backend == "rust")
+    ocr_server.expected_requests = 0
+    calls: Final = []
+
+    async def acquire() -> str:
+        calls.append("awaited")
+        return "unused"
+
+    coroutine: Final = acquire()
+
+    def provider() -> object:
+        return coroutine
+
+    try:
+        with pytest.raises(litellm.APIConnectionError, match="Azure AD token must be a string"):
+            await call_aocr(
+                ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider
+            )
+    finally:
+        coroutine.close()
+    assert calls == []
+    assert ocr_server.requests == []

tests/test_litellm_rust/support/__init__.py

@@ -0,0 +1 @@
+

tests/test_litellm_rust/support/callback_recorder.py

@@ -0,0 +1,178 @@
+import asyncio
+import copy
+import threading
+import time
+from dataclasses import dataclass
+from typing import Final
+
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
+
+
+async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None:
+    # The wrapper's scheduled helper enqueues without awaiting; let it run before joining the queue.
+    await asyncio.sleep(0)
+    worker.start()
+    await asyncio.wait_for(worker.flush(), timeout=10)
+
+
+@dataclass(frozen=True, slots=True)
+class HookEvent:
+    name: str
+    call_type: str | None
+    stream: bool | None
+    thread: threading.Thread
+    loop: asyncio.AbstractEventLoop | None
+    has_running_loop: bool
+    kwargs: object
+    response: object
+
+
+class RecordingLogger(CustomLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        self._events: list[HookEvent] = []
+        self._condition = threading.Condition()
+
+    @property
+    def events(self) -> tuple[HookEvent, ...]:
+        with self._condition:
+            return tuple(self._events)
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return tuple(event.name for event in self.events)
+
+    def _record(self, name: str, kwargs: object = None, response: object = None) -> None:
+        details: Final = kwargs if isinstance(kwargs, dict) else {}
+        try:
+            snapshot: Final = copy.deepcopy(details)
+        except Exception:
+            snapshot = dict(details)
+        if "exception" in details:
+            snapshot["exception"] = details["exception"]
+        try:
+            loop: Final = asyncio.get_running_loop()
+            has_running_loop: Final = True
+        except RuntimeError:
+            loop = None
+            has_running_loop = False
+        event: Final = HookEvent(
+            name=name,
+            call_type=details.get("call_type"),
+            stream=details.get("stream"),
+            thread=threading.current_thread(),
+            loop=loop,
+            has_running_loop=has_running_loop,
+            kwargs=snapshot,
+            response=response,
+        )
+        with self._condition:
+            self._events.append(event)
+            self._condition.notify_all()
+
+    def wait_for(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]:
+        deadline: Final = time.monotonic() + timeout
+        with self._condition:
+            while sum(event.name == name for event in self._events) < count:
+                remaining: Final = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise TimeoutError(f"Timed out waiting for {count} {name} events; saw {self.names}")
+                self._condition.wait(remaining)
+            return tuple(event for event in self._events if event.name == name)
+
+    async def wait_for_async(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]:
+        await asyncio.wait_for(asyncio.to_thread(self.wait_for, name, count, timeout), timeout=timeout + 1)
+        from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+        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):
+        self._record("log_pre_api_call", kwargs)
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("log_success_event", kwargs, response_obj)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("async_log_success_event", kwargs, response_obj)
+
+    def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("log_stream_event", kwargs, response_obj)
+
+    async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("async_log_stream_event", kwargs, response_obj)
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("log_failure_event", kwargs, response_obj)
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._record("async_log_failure_event", kwargs, response_obj)
+
+    def logging_hook(self, kwargs, result, call_type):
+        self._record("logging_hook", kwargs, result)
+        return kwargs, result
+
+    async def async_logging_hook(self, kwargs, result, call_type):
+        self._record("async_logging_hook", kwargs, result)
+        return kwargs, result
+
+    async def async_pre_call_deployment_hook(self, kwargs, call_type):
+        self._record("async_pre_call_deployment_hook", kwargs)
+
+    async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+        self._record("async_post_call_success_deployment_hook", request_data, response)
+        return response
+
+    async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, fallback_depth=None):
+        self._record("async_post_call_failure_deployment_hook", request_data, exception)
+
+
+@dataclass(frozen=True, slots=True)
+class LiveReferenceEvent:
+    kwargs: object
+    response: object
+
+
+class LiveReferenceLogger(CustomLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        self._events: list[LiveReferenceEvent] = []
+        self._condition = threading.Condition()
+
+    @property
+    def events(self) -> tuple[LiveReferenceEvent, ...]:
+        with self._condition:
+            return tuple(self._events)
+
+    def _record(self, kwargs: object, response: object) -> None:
+        with self._condition:
+            self._events.append(LiveReferenceEvent(kwargs=kwargs, response=response))
+            self._condition.notify_all()
+
+    def wait_for(self, count: int = 1, timeout: float = 10) -> tuple[LiveReferenceEvent, ...]:
+        deadline: Final = time.monotonic() + timeout
+        with self._condition:
+            while len(self._events) < count:
+                remaining: Final = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise TimeoutError(f"Timed out waiting for {count} live-reference events")
+                self._condition.wait(remaining)
+            return tuple(self._events)
+
+    async def wait_for_async(self, count: int = 1, timeout: float = 10) -> tuple[LiveReferenceEvent, ...]:
+        return await asyncio.wait_for(asyncio.to_thread(self.wait_for, count, timeout), timeout=timeout + 1)
+
+    def release(self) -> None:
+        with self._condition:
+            self._events.clear()
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._record(kwargs, response_obj)
+
+
+class SecondaryLiveReferenceLogger(LiveReferenceLogger):
+    pass

tests/test_litellm_rust/support/prometheus.py

@@ -0,0 +1,19 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Final
+
+from prometheus_client import REGISTRY, CollectorRegistry
+
+
+@contextmanager
+def isolated_prometheus_registry(registry: CollectorRegistry = REGISTRY) -> Iterator[None]:
+    original: Final = tuple(registry._collector_to_names)  # pyright: ignore[reportPrivateUsage]  # prometheus_client has no public collector enumeration
+    for collector in original:
+        registry.unregister(collector)
+    try:
+        yield
+    finally:
+        for collector in tuple(registry._collector_to_names):  # pyright: ignore[reportPrivateUsage]  # remove collectors created inside the isolation scope
+            registry.unregister(collector)
+        for collector in original:
+            registry.register(collector)

tests/test_litellm_rust/support/provenance.py

@@ -0,0 +1,4 @@
+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/support/recording_server.py

@@ -0,0 +1,145 @@
+import asyncio
+import copy
+import json
+import threading
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Final
+
+import pytest
+
+
+@dataclass
+class RecordedRequest:
+    method: str
+    path: str
+    headers: dict[str, str]
+    raw_body: bytes
+    body: object | None
+
+
+@dataclass
+class ResponseSpec:
+    body: object
+    status: int = 200
+    headers: dict[str, str] = field(default_factory=dict)
+    delay: float = 0
+    events: tuple[tuple[str, object], ...] = ()
+    chunks: tuple[bytes, ...] = ()
+    accepted: threading.Event | None = None
+    release_before_response: threading.Event | None = None
+    release: threading.Event | None = None
+
+
+@dataclass
+class RecordingServer:
+    server: ThreadingHTTPServer
+    requests: list[RecordedRequest]
+    responses: list[ResponseSpec]
+    default_response: ResponseSpec
+    expected_requests: int | None = 1
+
+    @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)
+
+    async def wait_for_requests(self, count: int) -> None:
+        async with asyncio.timeout(2):
+            while len(self.requests) < count:
+                await asyncio.sleep(0.01)
+
+
+@contextmanager
+def recording_service() -> Iterator[RecordingServer]:
+    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()},
+                    raw_body=raw_body,
+                    body=body,
+                )
+            )
+            response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response)
+            if response.accepted is not None:
+                response.accepted.set()
+            if response.release_before_response is not None:
+                response.release_before_response.wait(timeout=10)
+            if response.delay:
+                time.sleep(response.delay)
+            payload: Final = (
+                b"".join(response.chunks)
+                if response.chunks
+                else (
+                    b"".join(
+                        f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in response.events
+                    )
+                    if response.events
+                    else json.dumps(response.body).encode()
+                )
+            )
+            self.send_response(response.status)
+            self.send_header(
+                "Content-Type", "text/event-stream" if response.events or response.chunks else "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:
+                if response.chunks:
+                    for index, chunk in enumerate(response.chunks):
+                        if index == 1 and response.release is not None:
+                            response.release.wait(timeout=10)
+                        self.wfile.write(chunk)
+                        self.wfile.flush()
+                else:
+                    self.wfile.write(payload)
+            except (BrokenPipeError, ConnectionResetError):
+                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:
+        recording_server = RecordingServer(
+            server=server,
+            requests=requests,
+            responses=responses,
+            default_response=ResponseSpec(body={}),
+        )
+        yield recording_server
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()
+        if recording_server.expected_requests is not None:
+            assert len(recording_server.requests) == recording_server.expected_requests
+        assert recording_server.responses == []
+
+
+@pytest.fixture
+def recording_server() -> Iterator[RecordingServer]:
+    with recording_service() as server:
+        yield server

tests/test_litellm_rust/support/requests.py

@@ -0,0 +1,113 @@
+from typing import Final
+
+import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from litellm.rust_bridge import ocr as native_ocr
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
+OCR_MODEL: Final = "mistral/mistral-ocr-latest"
+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},
+}
+
+MESSAGES_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929"
+MESSAGES: Final = [{"role": "user", "content": "Hello"}]
+MESSAGES_RESPONSE: Final = {
+    "id": "msg_native",
+    "type": "message",
+    "role": "assistant",
+    "model": "claude-sonnet-4-5-20250929",
+    "content": [{"type": "text", "text": "Hello from native Messages"}],
+    "stop_reason": "end_turn",
+    "stop_sequence": None,
+    "usage": {"input_tokens": 5, "output_tokens": 4},
+}
+CHAT_MODEL: Final = "anthropic/claude-opus-5"
+CHAT_MESSAGES: Final = [
+    {"role": "system", "content": "Keep the answer short"},
+    {"role": "user", "content": "Earlier safe question"},
+    {"role": "assistant", "content": "Earlier safe answer"},
+    {"role": "user", "content": "Employee SSN: 078-05-1120"},
+]
+CHAT_RESPONSE: Final = {
+    "id": "msg_chat_native",
+    "type": "message",
+    "role": "assistant",
+    "model": "claude-opus-5",
+    "content": [{"type": "text", "text": "Handled safely"}],
+    "stop_reason": "end_turn",
+    "stop_sequence": None,
+    "usage": {"input_tokens": 8, "output_tokens": 2},
+}
+
+
+def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
+    return {
+        "model": OCR_MODEL,
+        "document": dict(OCR_DOCUMENT),
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        **kwargs,
+    }
+
+
+def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
+    if not isinstance(response, OCRResponse):
+        raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
+    return response
+
+
+async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    return await litellm.aocr(**ocr_arguments(server, **kwargs))
+
+
+def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    return native_ocr.ocr(ocr_arguments(server, **kwargs))
+
+
+async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    return await native_ocr.aocr(ocr_arguments(server, **kwargs))
+
+
+def request_body(kwargs: dict[str, object]) -> dict[str, object]:
+    additional_args = kwargs["additional_args"]
+    assert isinstance(additional_args, dict)
+    body = additional_args["complete_input_dict"]
+    assert isinstance(body, dict)
+    return body
+
+
+def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
+    additional_args = kwargs["additional_args"]
+    assert isinstance(additional_args, dict)
+    headers = additional_args["headers"]
+    assert isinstance(headers, dict)
+    return headers
+
+
+MESSAGES_EVENTS: Final = (
+    ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}),
+    ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
+    (
+        "content_block_delta",
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "text_delta", "text": "Hello from native Messages"},
+        },
+    ),
+    ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+    (
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"input_tokens": 5, "output_tokens": 4},
+        },
+    ),
+    ("message_stop", {"type": "message_stop"}),
+)

tests/test_litellm_rust/support/routes.py

@@ -0,0 +1,135 @@
+import asyncio
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from typing import Final, Literal
+
+import litellm
+from tests.test_litellm_rust.support.callback_recorder import HookEvent, RecordingLogger
+from tests.test_litellm_rust.support.requests import (
+    MESSAGES,
+    MESSAGES_EVENTS,
+    MESSAGES_MODEL,
+    MESSAGES_RESPONSE,
+    OCR_RESPONSE,
+    call_aocr,
+    call_ocr,
+)
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+RouteName = Literal["ocr-sync", "ocr-async", "messages", "messages-stream"]
+
+
+@dataclass(frozen=True, slots=True)
+class Route:
+    name: RouteName
+    call_type: str
+    provider_model: str
+    provider_response: Mapping[str, object]
+    expected_text: str
+    expected_cost: float
+    fires_async_hooks: bool
+
+    async def invoke(self, server: RecordingServer, **kwargs: object) -> object:
+        match self.name:
+            case "ocr-sync":
+                return await asyncio.to_thread(call_ocr, server, **kwargs)
+            case "ocr-async":
+                return await call_aocr(server, **kwargs)
+            case "messages":
+                return await _call_messages(server, **kwargs)
+            case "messages-stream":
+                stream: Final = await self.open_stream(server, **kwargs)
+                return [chunk async for chunk in stream]
+
+    async def open_stream(self, server: RecordingServer, **kwargs: object) -> AsyncIterator[object]:
+        if self.name != "messages-stream":
+            raise ValueError(f"{self.name} is not a streaming route")
+        stream: Final = await _call_messages(server, stream=True, **kwargs)
+        if not isinstance(stream, AsyncIterator):
+            raise TypeError(f"Expected async stream, got {type(stream).__name__}")
+        return stream
+
+    def response_text(self, response: object) -> str:
+        match self.name:
+            case "ocr-sync" | "ocr-async":
+                pages: Final = getattr(response, "pages", None)
+                if not isinstance(pages, list) or not pages:
+                    raise TypeError(f"Expected OCR response pages, got {type(response).__name__}")
+                markdown: Final = getattr(pages[0], "markdown", None)
+                if not isinstance(markdown, str):
+                    raise TypeError(f"Expected OCR markdown, got {type(markdown).__name__}")
+                return markdown
+            case "messages":
+                if not isinstance(response, Mapping):
+                    raise TypeError(f"Expected Messages mapping, got {type(response).__name__}")
+                content: Final = response.get("content")
+                if not isinstance(content, list) or not content or not isinstance(content[0], Mapping):
+                    raise TypeError("Expected Messages text content")
+                text: Final = content[0].get("text")
+                if not isinstance(text, str):
+                    raise TypeError(f"Expected Messages text, got {type(text).__name__}")
+                return text
+            case "messages-stream":
+                raise ValueError("Streaming response text is assembled by the stream consumer")
+
+
+async def _call_messages(server: RecordingServer, **kwargs: object) -> object:
+    return await litellm.anthropic.messages.acreate(
+        model=MESSAGES_MODEL,
+        messages=MESSAGES,
+        max_tokens=64,
+        api_key="test-key",
+        api_base=server.base_url,
+        **kwargs,
+    )
+
+
+MESSAGES_COST: Final = 5 * 3e-06 + 4 * 1.5e-05
+OCR_COST: Final = 0.004
+
+OCR_SYNC: Final = Route("ocr-sync", "ocr", "mistral-ocr-latest", OCR_RESPONSE, "native OCR response", OCR_COST, False)
+OCR_ASYNC: Final = Route("ocr-async", "aocr", "mistral-ocr-latest", OCR_RESPONSE, "native OCR response", OCR_COST, True)
+MESSAGES_ROUTE: Final = Route(
+    "messages",
+    "anthropic_messages",
+    "claude-sonnet-4-5-20250929",
+    MESSAGES_RESPONSE,
+    "Hello from native Messages",
+    MESSAGES_COST,
+    True,
+)
+MESSAGES_STREAM: Final = Route(
+    "messages-stream",
+    "anthropic_messages",
+    "claude-sonnet-4-5-20250929",
+    MESSAGES_RESPONSE,
+    "Hello from native Messages",
+    MESSAGES_COST,
+    True,
+)
+ALL_ROUTES: Final = (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE, MESSAGES_STREAM)
+ASYNC_ROUTES: Final = tuple(route for route in ALL_ROUTES if route.fires_async_hooks)
+NON_STREAM_ASYNC_ROUTES: Final = (OCR_ASYNC, MESSAGES_ROUTE)
+
+
+def route_id(route: Route) -> str:
+    return route.name
+
+
+def provider_response(route: Route) -> ResponseSpec:
+    return ResponseSpec(
+        body=route.provider_response,
+        events=MESSAGES_EVENTS if route.name == "messages-stream" else (),
+    )
+
+
+async def wait_for_callback(
+    route: Route,
+    recorder: RecordingLogger,
+    outcome: Literal["success", "failure"] = "success",
+    count: int = 1,
+) -> tuple[HookEvent, ...]:
+    event: Final = f"{'async_' if route.fires_async_hooks else ''}log_{outcome}_event"
+    if route.fires_async_hooks:
+        return await recorder.wait_for_async(event, count=count)
+    return recorder.wait_for(event, count=count)

tests/test_litellm_rust/test_provenance.py

@@ -0,0 +1,20 @@
+from typing import Final
+
+import pytest
+
+from litellm.rust_bridge.loader import get_native_bridge
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+def test_native_route_catalogue_matches_extension_exports() -> None:
+    from litellm.rust_bridge.provenance import (
+        RUST_NON_RESPONSE_ENTRYPOINTS,
+        RUST_RESPONSE_ENTRYPOINTS,
+    )
+
+    native: Final = get_native_bridge()
+    assert native is not None
+    route_exports: Final = frozenset(name for name in dir(native) if not name.startswith("_") and name[0].islower())
+
+    assert route_exports == RUST_RESPONSE_ENTRYPOINTS | RUST_NON_RESPONSE_ENTRYPOINTS

tests/test_litellm_rust/test_support.py

@@ -0,0 +1,115 @@
+import asyncio
+from pathlib import Path
+from types import ModuleType
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge import configuration
+from tests.test_litellm_rust import conftest
+from tests.test_litellm_rust.conftest import _isolated_list  # pyright: ignore[reportPrivateUsage]  # exercise fixture restoration without mutating SDK state
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("override", (None, False, True))
+@pytest.mark.parametrize("exception_type", (RuntimeError, asyncio.CancelledError))
+async def test_backend_restores_override_after_nested_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    override: bool | None,
+    exception_type: type[RuntimeError] | type[asyncio.CancelledError],
+) -> None:
+    monkeypatch.setenv("LITELLM_RUST", "1")
+    configuration.reset_rust_configuration()
+    if override is not None:
+        litellm.rust(override)
+
+    async def fail_in_nested_backend() -> None:
+        async with conftest.isolated_backend("python"):
+            assert configuration.rust_enabled() is False
+            with pytest.raises(exception_type, match="scope failed"):
+                await fail_in_rust_backend()
+            assert configuration.rust_enabled() is False
+            raise exception_type("scope failed")
+
+    async def fail_in_rust_backend() -> None:
+        async with conftest.isolated_backend("rust"):
+            assert configuration.rust_enabled() is True
+            raise exception_type("scope failed")
+
+    with pytest.raises(exception_type, match="scope failed"):
+        await fail_in_nested_backend()
+
+    assert configuration.rust_enabled() is (True if override is None else override)
+    monkeypatch.setenv("LITELLM_RUST", "0")
+    assert configuration.rust_enabled() is (False if override is None else override)
+
+
+@pytest.mark.asyncio
+async def test_backend_rejects_overlapping_tasks_without_changing_state() -> None:
+    async def competing_scope() -> None:
+        async with conftest.isolated_backend("rust"):
+            pytest.fail("overlapping backend scope was accepted")
+
+    async with conftest.isolated_backend("python"):
+        callback: Final = object()
+        litellm.callbacks.append(callback)
+        with pytest.raises(RuntimeError, match="isolated_backend scopes cannot overlap across tasks"):
+            await asyncio.create_task(competing_scope())
+        assert configuration.rust_enabled() is False
+        assert litellm.callbacks == [callback]
+        async with conftest.isolated_backend("rust"):
+            assert configuration.rust_enabled() is True
+            assert litellm.callbacks == []
+        assert configuration.rust_enabled() is False
+        assert litellm.callbacks == [callback]
+        with pytest.raises(RuntimeError, match="isolated_backend scopes cannot overlap across tasks"):
+            await asyncio.create_task(competing_scope())
+
+    async def subsequent_scope() -> None:
+        async with conftest.isolated_backend("rust"):
+            assert configuration.rust_enabled() is True
+
+    await asyncio.create_task(subsequent_scope())
+
+
+def test_callback_list_restores_identity_after_rebinding_and_failure() -> None:
+    container: Final = ModuleType("callback_registry")
+    callback: Final = object()
+    original: Final = [callback]
+    setattr(container, "callbacks", original)
+
+    def failing_test() -> None:
+        with _isolated_list(container, "callbacks"):
+            assert getattr(container, "callbacks") is original
+            assert original == []
+            original.append(object())
+            setattr(container, "callbacks", [object()])
+            raise RuntimeError("test failed")
+
+    with pytest.raises(RuntimeError, match="test failed"):
+        failing_test()
+
+    assert getattr(container, "callbacks") is original
+    assert original == [callback]
+
+
+def test_rust_extension_gate_skips_only_rust_suite_items(monkeypatch: pytest.MonkeyPatch) -> None:
+    class Item:
+        def __init__(self, path: Path) -> None:
+            self.path = path
+            self.markers: list[object] = []
+
+        def add_marker(self, marker: object) -> None:
+            self.markers.append(marker)
+
+    monkeypatch.setattr(conftest, "_parse_env_bool", lambda value: False)
+    rust_item: Final = Item(Path("tests/test_litellm_rust/ocr/test_dispatch.py"))
+    other_item: Final = Item(Path("tests/test_litellm/test_completion.py"))
+
+    conftest.pytest_collection_modifyitems([rust_item, other_item])  # pyright: ignore[reportArgumentType]  # lightweight items isolate hook selection
+
+    assert rust_item.markers
+    assert not other_item.markers