test(ocr): clarify retained Rust contracts

.github/workflows/test-rust.yml

@@ -6,11 +6,9 @@ on:
       - "litellm-rust/**"
       - "litellm/rust_bridge/**"
       - "tests/test_litellm_rust/**"
-      - "litellm/integrations/**"
+      - "litellm/integrations/custom_logger.py"
       - "litellm/litellm_core_utils/litellm_logging.py"
       - "litellm/litellm_core_utils/logging_worker.py"
-      - "litellm/proxy/guardrails/**"
-      - "litellm/proxy/common_request_processing.py"
       - "litellm/utils.py"
       - "litellm/ocr/**"
       - "litellm/llms/base_llm/ocr/**"
@@ -37,11 +35,9 @@ on:
       - "litellm-rust/**"
       - "litellm/rust_bridge/**"
       - "tests/test_litellm_rust/**"
-      - "litellm/integrations/**"
+      - "litellm/integrations/custom_logger.py"
       - "litellm/litellm_core_utils/litellm_logging.py"
       - "litellm/litellm_core_utils/logging_worker.py"
-      - "litellm/proxy/guardrails/**"
-      - "litellm/proxy/common_request_processing.py"
       - "litellm/utils.py"
       - "litellm/ocr/**"
       - "litellm/llms/base_llm/ocr/**"

tests/_prometheus_helpers.py

@@ -1,24 +0,0 @@
-from collections.abc import Iterator
-from contextlib import contextmanager
-from typing import Final
-
-from prometheus_client import REGISTRY, CollectorRegistry
-
-
-def clear_prometheus_registry() -> None:
-    for collector in tuple(REGISTRY._collector_to_names):  # pyright: ignore[reportPrivateUsage]  # prometheus_client has no public collector enumeration
-        REGISTRY.unregister(collector)
-
-
-@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/README.md

@@ -1,15 +1,13 @@
-# Rust bridge tests
+# Rust OCR bridge tests
 
-Put API-specific tests under `ocr/`, `messages/`, or `chat/`. Keep request construction, recording servers, callback recorders, and state isolation in `support/`. Under `integrations/`, keep registry inventory in `test_catalogue.py`, Python/Rust comparisons in `test_backend_parity.py`, callback ownership and completion in `test_callback_lifecycle.py`, individual exporters in `test_exporters.py`, and only multi-integration interactions in `test_composition.py`
+This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR behavior tests live under `ocr/`; reusable OCR request, callback, and recording-server fixtures live under `support/`
 
-Use `backend` with `"python"` and `"rust"` for parity tests. A test name must identify the observer and expected result. Keep multiple assertions together only when they are consequences of one mutation, failure, scheduling, or lifecycle event. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions. Shared route parametrization belongs in `support/routes.py` only when its observable contract is the same for every route
+A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions
 
-The catalogue stores labels for selected behavioral cases. Those labels are inventory metadata and do not prove that matching pytest cases exist or execute
+`ocr/test_requests.py` covers provider payloads, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_dispatch.py` covers public sync and async native dispatch, explicit Python dispatch, fallback, and the native compression header. `test_ocr.py` is the strict wire-level smoke test
 
 Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports
 
-Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture selects Rust unless a parity or fallback case explicitly selects Python. Backend selection alone does not prove native execution because a public route can fall back. Tests that claim native request coverage must also assert either the `x-litellm-rust` response marker or a wire property that distinguishes the native client. `test_public_ocr_executes_the_compiled_extension_without_python_fallback` is strict and its server rejects Python HTTPX requests
+Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture selects Rust for every test unless a fallback or parity case explicitly selects Python. Backend selection alone does not prove native execution because a public OCR request can fall back. Tests that claim native dispatch assert either the Rust response marker or a wire property that distinguishes the native client. `test_public_ocr_executes_the_compiled_extension_without_python_fallback` is strict and its server rejects Python HTTPX requests
 
-The retained callback contract modules are marked as non-strict expected failures until the implementation from #40070 lands. Harness tests and the compiled-extension OCR smoke test remain strict. Passing contract cases appear as XPASS so staging coverage stays 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
+The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The compiled-extension OCR smoke test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible

tests/test_litellm_rust/chat/__init__.py

@@ -1 +0,0 @@
-

tests/test_litellm_rust/chat/test_callback_mutation.py

@@ -1,470 +0,0 @@
-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], ids=["sync", "async"])
-@pytest.mark.parametrize("rebind_logging_view", [False, True], ids=["in-place-view", "rebound-view"])
-@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
-async def test_anthropic_pre_call_body_and_header_mutations_reach_provider(
-    recording_server: RecordingServer,
-    asynchronous: bool,
-    rebind_logging_view: bool,
-    native: bool,
-) -> None:
-    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(body=MESSAGES_RESPONSE)
-    callback_received_dict: 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"]
-            callback_received_dict.append(isinstance(body, dict))
-            if not isinstance(body, dict):
-                return
-            body["messages"][0]["content"][0]["text"] = "edited by callback"
-            body["max_tokens"] = 32
-            headers["x-retained-callback"] = "original"
-            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",
-        "messages": [{"role": "user", "content": "original"}],
-        "max_tokens": 64,
-        "api_key": "test-key",
-        "api_base": recording_server.base_url,
-        "callbacks": [recorder],
-        "num_retries": 0,
-    }
-    response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
-
-    assert response._hidden_params.get("additional_headers", {}).get("x-litellm-rust") == ("true" if native else None)
-    assert callback_received_dict == [True]
-    assert recorder.names.count("log_pre_api_call") == 1
-    assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == "edited by callback"
-    assert recording_server.requests[0].headers["x-retained-callback"] == "original"
-    assert recording_server.requests[0].body["max_tokens"] == 32
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("rebind_logging_view", [False, True], ids=["in-place-view", "rebound-view"])
-@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
-async def test_bedrock_pre_call_rebinding_does_not_replace_signed_transport(
-    recording_server: RecordingServer,
-    asynchronous: bool,
-    rebind_logging_view: bool,
-    native: bool,
-) -> None:
-    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-
-    litellm.rust(native)
-    recording_server.default_response = ResponseSpec(
-        body={
-            "output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}},
-            "stopReason": "end_turn",
-            "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
-            "metrics": {"latencyMs": 1},
-        }
-    )
-    callback_body_text: Final = []
-
-    class EditingLogger(RecordingLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            super().log_pre_api_call(model, messages, kwargs)
-            additional: Final = kwargs["additional_args"]
-            body: Final = additional["complete_input_dict"]
-            callback_body_text.append(
-                json.loads(body)["messages"][0]["content"][0]["text"] if isinstance(body, str) else None
-            )
-            additional["headers"]["x-retained-callback"] = "original"
-            if rebind_logging_view:
-                additional["complete_input_dict"] = {"replacement": True}
-                additional["headers"] = {"x-retained-callback": "replacement"}
-
-    recorder: Final = EditingLogger()
-    request: Final = {
-        "model": "bedrock/anthropic.claude-opus-5",
-        "messages": [{"role": "user", "content": "original"}],
-        "max_tokens": 64,
-        "api_base": recording_server.base_url,
-        "callbacks": [recorder],
-        "num_retries": 0,
-        "aws_access_key_id": "test-access",
-        "aws_secret_access_key": "test-secret",
-        "aws_region_name": "us-east-1",
-    }
-    response: Final = await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
-    recorded: Final = recording_server.requests[0]
-
-    assert response._hidden_params.get("additional_headers", {}).get("x-litellm-rust") == ("true" if native else None)
-    assert callback_body_text == ["original"]
-    assert recorder.names.count("log_pre_api_call") == 1
-    assert recorded.body["messages"][0]["content"][0]["text"] == "original"
-    assert "replacement" not in recorded.body
-    assert recorded.headers["x-retained-callback"] == "original"
-    assert recorded.body["inferenceConfig"]["maxTokens"] == 64
-    assert recorded.headers["authorization"].startswith("AWS4-HMAC-SHA256 ")
-    _verify_sigv4(recorded, "test-secret")
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("provider", ["anthropic", "bedrock"])
-@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
-async def test_chat_provider_failure_reaches_one_failure_callback(
-    recording_server: RecordingServer,
-    asynchronous: bool,
-    provider: str,
-    native: bool,
-) -> None:
-    from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-
-    litellm.rust(native)
-    recording_server.default_response = ResponseSpec(status=429, body={"message": "rate limited"})
-    recorder: Final = RecordingLogger()
-    request: 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 {}
-        ),
-    }
-
-    with pytest.raises((litellm.APIError, litellm.RateLimitError)) as raised:
-        await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
-    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 == 429
-    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
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("status", [200, 429])
-@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
-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]]] = []
-    shared_state: Final[list[tuple[object, object, object]]] = []
-    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)))
-            shared_state.append(
-                (
-                    isinstance(callback_body, str),
-                    json.loads(callback_body)["messages"][0]["content"][0]["text"]
-                    if isinstance(callback_body, str)
-                    else None,
-                    None,
-                )
-            )
-            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)))
-            shared_state.append(
-                (
-                    kwargs.get("callback_marker"),
-                    additional["complete_input_dict"] is replacement_body,
-                    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]
-    assert shared_state == [(True, "original", None), ("visible", True, True)]
-    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_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) -> object:
-        return 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()
-    responses: Final = 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)
-    assert all(response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" for response in responses)
-
-
-@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_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,17 +1,13 @@
 import asyncio
 import os
-from collections.abc import AsyncIterator, Generator
+from collections.abc import AsyncIterator, Generator, Iterator
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import ExitStack, asynccontextmanager, contextmanager
-from threading import Lock
+from contextlib import ExitStack, contextmanager
 from types import ModuleType
-from typing import Final, Literal, cast
+from typing import Final, 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
@@ -21,7 +17,6 @@ from litellm.rust_bridge.configuration import (  # pyright: ignore[reportPrivate
     _CONFIGURATION,
     _parse_env_bool,
 )
-from tests._prometheus_helpers import isolated_prometheus_registry
 from tests.test_litellm_rust.support.callback_recorder import drain_logging
 from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
 
@@ -34,46 +29,11 @@ CALLBACK_ATTRIBUTES: Final = (
     "_async_success_callback",
     "_async_failure_callback",
 )
-Backend = Literal["python", "rust"]
-EXPECTED_FAILURE_FILES: Final = frozenset(
-    {
-        "chat/test_callback_mutation.py",
-        "integrations/test_backend_parity.py",
-        "integrations/test_callback_lifecycle.py",
-        "integrations/test_composition.py",
-        "integrations/test_exporters.py",
-        "integrations/test_guardrails.py",
-        "messages/test_callback_mutation.py",
-        "messages/test_streaming.py",
-        "ocr/test_callbacks.py",
-        "ocr/test_dispatch.py",
-        "ocr/test_requests.py",
-        "test_provenance.py",
-    }
-)
-
-
-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()
+EXPECTED_FAILURE_REASONS: Final = {
+    "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070",
+    "ocr/test_dispatch.py": "requires the OCR native dispatch implementation from #40070",
+    "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070",
+}
 
 
 def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
@@ -84,7 +44,7 @@ def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
 
 
 @contextmanager
-def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]:
+def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]:
     source: Final = _list_attribute(container, attribute)
     original: Final = list(source)
     source.clear()  # mutable-ok: test isolation mutates global registries by design
@@ -97,7 +57,7 @@ def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]:
 
 
 @contextmanager
-def _rebound(container: object, attribute: str, value: object) -> Generator[None]:
+def _rebound(container: object, attribute: str, value: object) -> Iterator[None]:
     original: Final[object] = getattr(container, attribute)
     setattr(container, attribute, value)
     try:
@@ -107,36 +67,24 @@ def _rebound(container: object, attribute: str, value: object) -> Generator[None
 
 
 @contextmanager
-def _rust_mode(enabled: bool) -> Generator[None]:
+def _rust_mode(enabled: bool) -> Iterator[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]:
+@pytest_asyncio.fixture(autouse=True, loop_scope="function")
+async def isolate_ocr_test_state() -> AsyncIterator[None]:
     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(_isolated_list(litellm_logging, "_in_memory_loggers"))  # pyright: ignore[reportPrivateUsage]  # no public callback-cache accessor
+        stack.enter_context(_rebound(utils, "callback_list", []))  # rebind-ok: isolate legacy callback registry
+        stack.enter_context(_rebound(litellm, "cache", None))  # test-quality-ok: isolate process-global cache
+        stack.enter_context(_rust_mode(True))
+        executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging")
         stack.enter_context(_rebound(utils, "executor", executor))
         try:
-            yield stack
+            yield
         finally:
             try:
                 await drain_logging()
@@ -145,13 +93,6 @@ async def _isolated_backend(backend: Backend) -> AsyncIterator[ExitStack]:
                 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:
@@ -159,16 +100,13 @@ def recording_server() -> Generator[RecordingServer]:
 
 
 def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
-    expected_failure: Final = pytest.mark.xfail(
-        reason="requires the retained callback implementation from #40070",
-        strict=False,
-    )
     for item in items:
         if "test_litellm_rust" not in item.path.parts:
             continue
         relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :])
-        if relative_path in EXPECTED_FAILURE_FILES:
-            item.add_marker(expected_failure)
+        reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path)
+        if reason is not None:
+            item.add_marker(pytest.mark.xfail(reason=reason, strict=False))
 
     if not _parse_env_bool(os.environ.get("LITELLM_RUST")):
         skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")

tests/test_litellm_rust/integrations/__init__.py

@@ -1,75 +0,0 @@
-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

@@ -1,193 +0,0 @@
-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_case_labels: 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_case_labels=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

@@ -1,43 +0,0 @@
-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

@@ -1,203 +0,0 @@
-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_backend_parity.py

@@ -1,92 +0,0 @@
-from typing import Final
-
-import pytest
-
-import litellm
-from tests.test_litellm_rust.conftest import Backend, isolated_backend
-from tests.test_litellm_rust.integrations import (
-    MESSAGES_ROUTE,
-    OCR_ASYNC,
-    OCR_SYNC,
-    Route,
-    RunObservation,
-    provider_response,
-    route_id,
-    wait_for_callback,
-)
-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.recording_server import ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import CHAT_MESSAGES, CHAT_MODEL, CHAT_RESPONSE
-
-pytestmark = pytest.mark.requires_rust_extension
-
-
-async def observe_route(backend: Backend, route: Route) -> 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),
-            )
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE), ids=route_id)
-async def test_python_and_rust_match_public_route_outputs_and_logging(route: Route) -> None:
-    python: Final = await observe_route("python", route)
-    rust: Final = await observe_route("rust", route)
-    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.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
-    assert python_body == rust_body
-    assert python.rust_dispatch is False
-    assert rust.rust_dispatch is True
-
-
-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_python_and_rust_emit_the_same_retry_callback_sequence() -> None:
-    python_names, python_used_rust = await observe_retry("python")
-    rust_names, rust_used_rust = 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_used_rust is False
-    assert rust_used_rust is True

tests/test_litellm_rust/integrations/test_callback_lifecycle.py

@@ -1,281 +0,0 @@
-import asyncio
-import gc
-import json
-import threading
-import weakref
-from collections.abc import Mapping
-from datetime import datetime
-from typing import Final
-
-import pytest
-
-import litellm
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.types.utils import CallTypes
-from tests.test_litellm_rust.conftest import Backend, isolated_backend
-from tests.test_litellm_rust.integrations import MESSAGES_ROUTE, MESSAGES_STREAM, OCR_ASYNC, Route, provider_response
-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 ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS
-
-pytestmark = pytest.mark.requires_rust_extension
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ("python", "rust"))
-async def test_terminal_callbacks_receive_the_same_live_kwargs_and_response(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_redacted_and_plain_loggers_receive_separate_payloads_with_shared_metadata(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)
-            response: Final = 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)
-            assert has_rust_response_marker(response) is (backend == "rust")
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ("python", "rust"))
-@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected"))
-async def test_deferred_terminal_callback_runs_only_after_proxy_completion_decision(
-    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_retains_metadata_until_completion(backend: Backend) -> None:
-    class Root:
-        pass
-
-    class Suspended(CustomLogger):
-        def __init__(self) -> None:
-            super().__init__()
-            self.started = asyncio.Event()
-            self.release = asyncio.Event()
-            self.saw_retained_root = False
-
-        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
-            self.saw_retained_root = 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)
-                assert callback.saw_retained_root is True
-                gc.collect()
-                assert reference() is not None
-                assert has_rust_response_marker(response) is (backend == "rust")
-                callback.release.set()
-                await drain_logging()
-                del response
-                return reference
-
-    reference: Final = await invoke()
-    await asyncio.sleep(0)
-    gc.collect()
-    assert reference() is None
-
-
-@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_terminal_failure_without_calling_deployment_failure_hook(
-    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("ending", ("provider_error", "truncated", "close"))
-async def test_established_rust_stream_finishes_once_without_deployment_failure_hook(ending: str) -> None:
-    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)
-            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_catalogue.py

@@ -1,54 +0,0 @@
-from typing import Final
-
-import pytest
-
-from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
-from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry
-from litellm.types.guardrails import SupportedGuardrailIntegrations
-from tests.test_litellm_rust.integrations import (
-    DISCOVERED_ONLY_GUARDRAIL_NAMES,
-    ENTERPRISE_LOGGER_NAMES,
-    GUARDRAIL_NAMES,
-    GUARDRAIL_OBLIGATIONS,
-    LOGGER_OBLIGATIONS,
-    OSS_LOGGER_NAMES,
-    REQUIRED_GUARDRAIL_BEHAVIOR,
-    REQUIRED_LOGGER_BEHAVIOR,
-)
-
-pytestmark = pytest.mark.requires_rust_extension
-
-
-def test_logger_catalogue_matches_registered_integration_names() -> None:
-    registered_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 <= registered_names
-    assert registered_names <= OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
-    assert catalogued_names == OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
-
-
-def test_guardrail_catalogue_matches_registered_integration_names() -> 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_case_labels() -> None:
-    labelled_loggers: Final = frozenset(
-        name
-        for obligation in LOGGER_OBLIGATIONS.values()
-        if obligation.behavioral_case_labels
-        for name in obligation.registration_names
-    )
-    labelled_guardrails: Final = frozenset(
-        name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_case_labels
-    )
-
-    assert REQUIRED_LOGGER_BEHAVIOR <= labelled_loggers
-    assert REQUIRED_GUARDRAIL_BEHAVIOR <= labelled_guardrails

tests/test_litellm_rust/integrations/test_composition.py

@@ -1,149 +0,0 @@
-import json
-from dataclasses import dataclass
-from typing import Final
-
-import pytest
-from opentelemetry.trace import StatusCode
-
-from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
-from litellm.integrations.prometheus import PrometheusLogger
-from tests.test_litellm_rust.conftest import Backend, isolated_backend
-from tests.test_litellm_rust.integrations import (
-    ASYNC_ROUTES,
-    MESSAGES_ROUTE,
-    OCR_ASYNC,
-    AsyncBoundaryLogger,
-    MutatingFailingLogger,
-    OtelHarness,
-    Route,
-    gcs_literalai_harness,
-    metric_value,
-    provider_response,
-    route_id,
-    wait_for_callback,
-)
-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, recording_service
-
-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"],
-        },
-    }
-]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=lambda route: f"otel-prometheus-generic-api-{route.name}-success")
-async def test_otel_prometheus_and_generic_api_each_export_one_success(
-    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")
-        response: Final = 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
-        assert has_rust_response_marker(response)
-
-
-@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"
-        assert has_rust_response_marker(response)
-
-
-@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")

tests/test_litellm_rust/integrations/test_exporters.py

@@ -1,253 +0,0 @@
-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._prometheus_helpers import isolated_prometheus_registry
-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_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)
-
-
-async def invoke_native(route: Route, provider: RecordingServer, **kwargs: object) -> object:
-    if route.name == "messages-stream":
-        stream: Final = await route.open_stream(provider, **kwargs)
-        assert has_rust_response_marker(stream)
-        return [chunk async for chunk in stream]
-    response: Final = await route.invoke(provider, **kwargs)
-    assert has_rust_response_marker(response)
-    return response
-
-
-@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 invoke_native(route, 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 invoke_native(route, 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
-async def test_otel_stream_span_appears_after_an_interrupted_consumer_closes(
-    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])
-    first_chunk: Final = await anext(stream)
-    await drain_logging()
-
-    assert otel.spans() == ()
-    assert has_rust_response_marker(stream)
-
-    await stream.aclose()
-    await drain_logging()
-
-    assert first_chunk is not None
-    assert len(otel.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_export_matching_standard_logging_values(
-    route: Route, provider: RecordingServer, otel: OtelHarness
-) -> None:
-    recorder: Final = RecordingLogger()
-    await invoke_native(route, 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 invoke_native(route, 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 invoke_native(MESSAGES_ROUTE, 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 invoke_native(route, provider, callbacks=[recorder])
-    await invoke_native(route, 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 invoke_native(OCR_SYNC, 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

@@ -1,697 +0,0 @@
-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,
-    )
-
-
-async def _observed_stream(source: AsyncIterator[object], first_chunk: asyncio.Event) -> AsyncIterator[object]:
-    async for chunk in source:
-        first_chunk.set()
-        yield chunk
-
-
-@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()
-            guarded = arm_guarded_stream(_observed_stream(source, first_chunk), 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

@@ -1 +0,0 @@
-

tests/test_litellm_rust/messages/test_callback_mutation.py

@@ -1,688 +0,0 @@
-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], ids=["response", "stream"])
-@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["edit-returns", "edit-raises"])
-@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
-async def test_messages_pre_call_edits_reach_later_callbacks_but_only_headers_reach_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"
-    assert has_rust_response_marker(response) is native
-
-
-@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], ids=["python", "rust"])
-@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"])
-@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["edit-returns", "edit-raises"])
-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 = []
-    aliases: Final = []
-
-    class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            body = request_body(kwargs)
-            message = kwargs["messages"][0]
-            aliases.append(
-                (
-                    body["messages"][0] is message,
-                    body["messages"][0]["content"] is message["content"],
-                    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 aliases == [(True, True, True)]
-    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"
-    assert has_rust_response_marker(response) is native
-
-
-@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

@@ -1,275 +0,0 @@
-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/test_callbacks.py

@@ -9,7 +9,6 @@ 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,
@@ -18,6 +17,7 @@ from tests.test_litellm_rust.support.requests import (
     request_body,
     request_headers,
 )
+from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
 
 pytestmark = pytest.mark.requires_rust_extension
@@ -37,19 +37,18 @@ async def call_aocr(server: RecordingServer, callbacks: list[CustomLogger], **kw
     return await call_native_aocr(server, callbacks=callbacks, **kwargs)
 
 
-def test_pre_call_receives_expected_provider_request(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_callback_receives_transformed_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"])))
+        def log_pre_api_call(self, model, _messages, kwargs):
+            observations.append((model, copy.deepcopy(kwargs["additional_args"])))
 
     call_ocr(ocr_server, [Observe()], pages=[0])
 
     assert len(observations) == 1
-    model, messages, additional_args = observations[0]
+    model, 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",
@@ -58,20 +57,20 @@ def test_pre_call_receives_expected_provider_request(ocr_server: RecordingServer
     }
 
 
-@pytest.mark.parametrize("raise_after_edit", [False, True])
-def test_pre_call_body_edits_reach_later_callbacks_and_provider(
+@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"])
+def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
     ocr_server: RecordingServer, raise_after_edit: bool
 ) -> None:
     observed: Final = []
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             request_body(kwargs)["include_image_base64"] = True
             if raise_after_edit:
-                raise RuntimeError("audit exporter unavailable")
+                raise RuntimeError("pre-call callback failed")
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             observed.append(copy.deepcopy(request_body(kwargs)))
 
     call_ocr(ocr_server, [Edit(), Observe()], include_image_base64=False)
@@ -80,15 +79,15 @@ def test_pre_call_body_edits_reach_later_callbacks_and_provider(
     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:
+def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None:
     observed: Final = []
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             request_headers(kwargs)["x-audit-tag"] = "reviewed"
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             observed.append(dict(request_headers(kwargs)))
 
     call_ocr(ocr_server, [Edit(), Observe()])
@@ -98,9 +97,9 @@ def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: Re
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("rust_enabled", [False, True], ids=["python", "rust"])
+@pytest.mark.parametrize("rust_enabled", [False, True], ids=["python-backend", "rust-backend"])
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_pre_call_nested_mutation_updates_retained_references(
+async def test_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
     ocr_server: RecordingServer, rust_enabled: bool, asynchronous: bool
 ) -> None:
     litellm.rust(rust_enabled)
@@ -110,12 +109,12 @@ async def test_pre_call_nested_mutation_updates_retained_references(
     aliases: Final = []
 
     class Retain(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             aliases.append(request_body(kwargs)["document"] is original)
             retained.append(request_body(kwargs)["document"])
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             original["document_url"] = replacement_url
 
     arguments: Final = {
@@ -134,13 +133,15 @@ async def test_pre_call_nested_mutation_updates_retained_references(
     assert has_rust_response_marker(response) is rust_enabled
 
 
-def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_document(
+    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):
+        def log_pre_api_call(self, model, _messages, kwargs):
             body = request_body(kwargs)
             retained.append(body["document"])
             body["document"] = replacement
@@ -156,15 +157,17 @@ def test_pre_call_field_replacement_preserves_original_references(ocr_server: Re
     assert ocr_server.requests[0].body["document"] == replacement
 
 
-def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider(
+    ocr_server: RecordingServer,
+) -> None:
     observed: Final = []
 
     class Rebind(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
 
     class Observe(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             observed.append(request_body(kwargs))
 
     call_ocr(ocr_server, [Rebind(), Observe()])
@@ -173,15 +176,15 @@ def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: R
     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:
+def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None:
     queued: Final = []
 
     class QueuePayload(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             queued.append(request_body(kwargs))
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             request_body(kwargs)["queued-edit"] = True
 
     call_ocr(ocr_server, [QueuePayload(), Edit()])
@@ -189,13 +192,13 @@ def test_queued_payload_observes_later_callback_mutations(ocr_server: RecordingS
     assert queued[0]["queued-edit"] is True
 
 
-def test_pre_call_state_reaches_terminal_callbacks(ocr_server: RecordingServer) -> None:
+def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(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):
+        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):
@@ -209,7 +212,9 @@ def test_pre_call_state_reaches_terminal_callbacks(ocr_server: RecordingServer)
 
 
 @pytest.mark.asyncio
-async def test_success_callbacks_receive_expected_context_and_response(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_success_callback_receives_call_id_metadata_and_response(
+    ocr_server: RecordingServer,
+) -> None:
     recorder: Final = RecordingLogger()
 
     await call_aocr(
@@ -228,7 +233,9 @@ async def test_success_callbacks_receive_expected_context_and_response(ocr_serve
 
 
 @pytest.mark.asyncio
-async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response(
+    ocr_server: RecordingServer,
+) -> None:
     ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
     observations: Final = []
 
@@ -249,7 +256,7 @@ async def test_failure_callbacks_receive_expected_context_and_error(ocr_server:
 
 
 @pytest.mark.asyncio
-async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None:
     caller_loop: Final = asyncio.get_running_loop()
     caller_thread: Final = threading.current_thread()
     recorder: Final = RecordingLogger()
@@ -263,13 +270,15 @@ async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingS
 
 
 @pytest.mark.asyncio
-async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback(
+    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):
+        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):
@@ -286,21 +295,21 @@ async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: Recordin
 
 
 @pytest.mark.asyncio
-async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_callbacks(
+async def test_native_aocr_callback_error_does_not_mask_provider_error_or_skip_later_failure_callbacks(
     ocr_server: RecordingServer,
 ) -> None:
     ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
     recorder: Final = RecordingLogger()
 
-    class UnavailableExporter(CustomLogger):
+    class FailingCallback(CustomLogger):
         def log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            raise RuntimeError("exporter unavailable")
+            raise RuntimeError("failure callback failed")
 
         async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            raise RuntimeError("exporter unavailable")
+            raise RuntimeError("failure callback failed")
 
     with pytest.raises(litellm.InternalServerError) as caught:
-        await call_aocr(ocr_server, [UnavailableExporter(), recorder])
+        await call_aocr(ocr_server, [FailingCallback(), 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")
@@ -311,7 +320,9 @@ async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_
     assert "async_log_success_event" not in recorder.names
 
 
-def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: RecordingServer) -> None:
+def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times(
+    ocr_server: RecordingServer,
+) -> None:
     recorder: Final = RecordingLogger()
 
     call_ocr(
@@ -329,8 +340,8 @@ def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: Recordi
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_azure_token_callback_precedes_logger_and_preserves_context(
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_public_azure_ocr_resolves_token_before_pre_call_on_caller_context(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
@@ -353,7 +364,7 @@ async def test_azure_token_callback_precedes_logger_and_preserves_context(
             return "caller-token"
 
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model, _messages, kwargs):
             assert request_headers(kwargs)["Authorization"] == "Bearer caller-token"
             observations.append("pre_call")
             request_headers(kwargs)["Authorization"] = "Bearer edited"
@@ -374,8 +385,8 @@ async def test_azure_token_callback_precedes_logger_and_preserves_context(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_azure_token_callback_can_reenter_native_sdk(
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_public_azure_ocr_token_provider_can_make_nested_native_ocr_call(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
@@ -408,7 +419,7 @@ async def test_azure_token_callback_can_reenter_native_sdk(
 
 
 @pytest.mark.asyncio
-async def test_concurrent_azure_token_callbacks_keep_results_and_errors_separate(
+async def test_concurrent_public_azure_ocr_calls_isolate_token_results_and_error(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
 ) -> None:
@@ -447,7 +458,7 @@ async def test_concurrent_azure_token_callbacks_keep_results_and_errors_separate
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"])
-async def test_azure_token_provider_is_released_after_request(
+async def test_public_azure_ocr_releases_token_provider_after_terminal_outcome(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     outcome: str,

tests/test_litellm_rust/ocr/test_dispatch.py

@@ -4,7 +4,7 @@ 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.response_marker 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
@@ -16,7 +16,7 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer:
     return recording_server
 
 
-def test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None:
+def test_sync_ocr_response_is_marked_as_rust_when_native_dispatch_is_enabled(ocr_server: RecordingServer) -> None:
     response: Final = litellm.ocr(
         model=OCR_MODEL,
         document=OCR_DOCUMENT,
@@ -29,7 +29,9 @@ def test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: Re
 
 
 @pytest.mark.asyncio
-async def test_public_aocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None:
+async def test_async_ocr_response_is_marked_as_rust_when_native_dispatch_is_enabled(
+    ocr_server: RecordingServer,
+) -> None:
     response: Final = await litellm.aocr(
         model=OCR_MODEL,
         document=OCR_DOCUMENT,
@@ -41,7 +43,7 @@ async def test_public_aocr_entrypoint_uses_native_transport_when_enabled(ocr_ser
     assert has_rust_response_marker(response)
 
 
-def test_public_ocr_falls_back_when_native_transport_declines(ocr_server: RecordingServer) -> None:
+def test_public_ocr_falls_back_to_python_for_unsupported_file_document(ocr_server: RecordingServer) -> None:
     response: Final = litellm.ocr(
         model=OCR_MODEL,
         document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
@@ -53,7 +55,9 @@ def test_public_ocr_falls_back_when_native_transport_declines(ocr_server: Record
     assert not has_rust_response_marker(response)
 
 
-def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingServer) -> None:
+def test_public_ocr_response_has_no_rust_marker_when_native_dispatch_is_disabled(
+    ocr_server: RecordingServer,
+) -> None:
     litellm.rust(False)
 
     response: Final = litellm.ocr(
@@ -67,7 +71,7 @@ def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingSer
     assert not has_rust_response_marker(response)
 
 
-def test_native_ocr_requests_an_uncompressed_response(ocr_server: RecordingServer) -> None:
+def test_native_ocr_sends_identity_accept_encoding_header(ocr_server: RecordingServer) -> None:
     litellm.ocr(
         model=OCR_MODEL,
         document=OCR_DOCUMENT,

tests/test_litellm_rust/ocr/test_requests.py

@@ -4,24 +4,23 @@ import pytest
 
 import litellm
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
 from tests.test_litellm_rust.support.requests import (
     OCR_DOCUMENT,
-    OCR_MODEL,
     OCR_RESPONSE,
-    call_native_ocr,
     call_aocr,
+    call_native_ocr,
     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
+from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_azure_ocr_calls_python_token_provider(
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_public_azure_ocr_uses_token_provider_result_as_bearer_token(
     ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
 ) -> None:
     calls: Final = []
@@ -61,7 +60,7 @@ def assert_native_request(server: RecordingServer) -> None:
     assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx")
 
 
-def test_ocr_sends_expected_provider_request(ocr_server: RecordingServer) -> None:
+def test_native_ocr_sends_model_and_document_to_mistral_ocr_path(ocr_server: RecordingServer) -> None:
     response: Final = call_ocr(ocr_server)
 
     assert response.pages[0].markdown == "native OCR response"
@@ -70,7 +69,7 @@ def test_ocr_sends_expected_provider_request(ocr_server: RecordingServer) -> Non
     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:
+def test_native_ocr_rejects_file_document_before_callbacks_or_provider_request(ocr_server: RecordingServer) -> None:
     ocr_server.expected_requests = 0
     recorder: Final = RecordingLogger()
 
@@ -85,21 +84,23 @@ def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: Reco
     assert recorder.events == ()
 
 
-def test_ocr_sends_optional_parameters(ocr_server: RecordingServer) -> None:
+def test_native_ocr_sends_pages_and_image_options(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:
+def test_native_ocr_merges_custom_headers_with_authorization(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:
+def test_native_mistral_ocr_uses_environment_api_key_when_argument_is_missing(
+    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
     monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
 
     call_native_ocr(ocr_server, api_key=None)
@@ -107,7 +108,7 @@ def test_ocr_resolves_provider_credentials(ocr_server: RecordingServer, monkeypa
     assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key"
 
 
-def test_ocr_explicit_credentials_override_defaults(
+def test_native_mistral_ocr_prefers_explicit_api_key_over_environment(
     ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
@@ -117,7 +118,9 @@ def test_ocr_explicit_credentials_override_defaults(
     assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
 
 
-def test_ocr_resolves_provider_endpoint(ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_native_azure_ocr_uses_environment_endpoint_and_api_key(
+    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)
 
@@ -128,7 +131,7 @@ def test_ocr_resolves_provider_endpoint(ocr_server: RecordingServer, monkeypatch
     assert ocr_server.requests[0].headers["authorization"] == "Bearer azure-key"
 
 
-def test_ocr_resolves_vertex_project_and_location(ocr_server: RecordingServer) -> None:
+def test_native_vertex_ocr_builds_path_from_project_and_location(ocr_server: RecordingServer) -> None:
     call_native_ocr(
         ocr_server,
         model="vertex_ai/mistral-ocr-2505",
@@ -143,7 +146,7 @@ def test_ocr_resolves_vertex_project_and_location(ocr_server: RecordingServer) -
     )
 
 
-def test_ocr_returns_normalized_response(ocr_server: RecordingServer) -> None:
+def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: RecordingServer) -> None:
     response: Final = call_ocr(ocr_server)
 
     assert isinstance(response, OCRResponse)
@@ -151,7 +154,7 @@ def test_ocr_returns_normalized_response(ocr_server: RecordingServer) -> None:
     assert response.usage_info.pages_processed == 1
 
 
-def test_ocr_provider_error_preserves_status_and_context(ocr_server: RecordingServer) -> None:
+def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: RecordingServer) -> None:
     ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
 
     with pytest.raises(litellm.BadRequestError) as caught:
@@ -163,7 +166,7 @@ def test_ocr_provider_error_preserves_status_and_context(ocr_server: RecordingSe
     assert "invalid OCR request" not in str(caught.value)
 
 
-def test_ocr_honors_request_timeout(ocr_server: RecordingServer) -> None:
+def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None:
     ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
 
     with pytest.raises(RuntimeError, match="OCR transport failed"):
@@ -173,8 +176,8 @@ def test_ocr_honors_request_timeout(ocr_server: RecordingServer) -> None:
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-@pytest.mark.parametrize("backend", ["python", "rust"])
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
 @pytest.mark.parametrize(
     "credentials, expected_token, expected_calls",
     [
@@ -182,8 +185,9 @@ def test_ocr_honors_request_timeout(ocr_server: RecordingServer) -> None:
         ({"azure_ad_token": "static-token"}, "callback-1", 1),
         ({"extra_headers": {"Authorization": "Bearer override"}}, "override", 1),
     ],
+    ids=["api-key-skips-provider", "provider-overrides-static-token", "header-overrides-provider"],
 )
-async def test_azure_ocr_token_provider_precedence(
+async def test_public_azure_ocr_applies_same_credential_precedence_on_python_and_rust(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
@@ -215,8 +219,8 @@ async def test_azure_ocr_token_provider_precedence(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_azure_ocr_does_not_cache_caller_tokens(
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_public_azure_ocr_calls_token_provider_for_each_request(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
@@ -250,10 +254,14 @@ class TokenAbort(BaseException):
 
 
 @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(
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
+@pytest.mark.parametrize(
+    "failure",
+    ["non_string", "type_error", "ordinary", "abort"],
+    ids=["non-string-result", "type-error", "value-error", "base-exception"],
+)
+async def test_public_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
@@ -307,8 +315,9 @@ async def test_azure_ocr_token_failure_stops_execution(
         {"model": "azure_ai/doc-intelligence/prebuilt-read"},
         {"document": {"type": "file", "file": b"pdf"}},
     ],
+    ids=["oidc-assertion", "document-intelligence-model", "file-document"],
 )
-def test_azure_unsupported_native_auth_does_not_invoke_provider(
+def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     configuration: dict[str, object],
@@ -336,8 +345,8 @@ def test_azure_unsupported_native_auth_does_not_invoke_provider(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ["python", "rust"])
-async def test_azure_missing_endpoint_prevents_token_callback(
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
+async def test_public_azure_ocr_validates_endpoint_before_calling_token_provider(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     backend: str,
@@ -363,8 +372,8 @@ async def test_azure_missing_endpoint_prevents_token_callback(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ["python", "rust"])
-async def test_azure_empty_callback_token_does_not_restore_static_token(
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
+async def test_public_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     backend: str,
@@ -387,8 +396,8 @@ async def test_azure_empty_callback_token_does_not_restore_static_token(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ["python", "rust"])
-async def test_azure_falsey_callable_leaves_static_token_unchanged(
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
+async def test_public_azure_ocr_ignores_falsey_token_provider_and_uses_static_token(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     backend: str,
@@ -417,8 +426,8 @@ async def test_azure_falsey_callable_leaves_static_token_unchanged(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("backend", ["python", "rust"])
-async def test_azure_token_callback_does_not_await_coroutine_result(
+@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
+async def test_public_azure_ocr_rejects_coroutine_returned_by_sync_token_provider(
     ocr_server: RecordingServer,
     isolated_azure_auth: None,
     backend: str,

tests/test_litellm_rust/support/callback_recorder.py

@@ -10,7 +10,6 @@ from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, Log
 
 
 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)
@@ -20,10 +19,8 @@ async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None:
 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
 
@@ -53,17 +50,13 @@ class RecordingLogger(CustomLogger):
             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,
         )
@@ -88,7 +81,7 @@ class RecordingLogger(CustomLogger):
         await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout)
         return tuple(event for event in self.events if event.name == name)
 
-    def log_pre_api_call(self, model, messages, kwargs):
+    def log_pre_api_call(self, model, _messages, kwargs):
         self._record("log_pre_api_call", kwargs)
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
@@ -97,12 +90,6 @@ class RecordingLogger(CustomLogger):
     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)
 
@@ -112,67 +99,3 @@ class RecordingLogger(CustomLogger):
     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/recording_server.py

@@ -9,8 +9,6 @@ from dataclasses import dataclass, field
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import Final
 
-import pytest
-
 
 @dataclass
 class RecordedRequest:
@@ -27,11 +25,6 @@ class ResponseSpec:
     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
@@ -76,44 +69,20 @@ def recording_service() -> Iterator[RecordingServer]:
                 )
             )
             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()
-                )
-            )
+            payload: Final = 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-Type", "application/json")
             self.send_header("Content-Length", str(len(payload)))
             for name, value in response.headers.items():
                 self.send_header(name, value)
             self.end_headers()
             try:
-                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)
+                self.wfile.write(payload)
             except (BrokenPipeError, ConnectionResetError):
                 pass
 
-        do_GET = _handle
         do_POST = _handle
 
         def log_message(self, format: str, *args: object) -> None:
@@ -137,9 +106,3 @@ def recording_service() -> Iterator[RecordingServer]:
         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

@@ -13,36 +13,6 @@ OCR_RESPONSE: Final = {
     "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 {
@@ -87,27 +57,3 @@ def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
     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

@@ -1,135 +0,0 @@
-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_ocr.py

@@ -1,6 +1,7 @@
 import json
 import threading
 from collections.abc import Generator
+from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import Final
 
@@ -11,17 +12,23 @@ import litellm
 pytestmark = pytest.mark.requires_rust_extension
 
 
+@dataclass(frozen=True, slots=True)
+class RecordedOCRRequest:
+    headers: dict[str, str]
+    body: object
+
+
 @pytest.fixture
-def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]]]:
-    requests: Final[list[dict[str, object]]] = []
+def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest]]]:
+    requests: Final[list[RecordedOCRRequest]] = []
 
     class Handler(BaseHTTPRequestHandler):
         def do_POST(self) -> None:
             requests.append(
-                {
-                    "headers": {name.lower(): value for name, value in self.headers.items()},
-                    "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
-                }
+                RecordedOCRRequest(
+                    headers={name.lower(): value for name, value in self.headers.items()},
+                    body=json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
+                )
             )
             if self.headers.get("User-Agent", "").startswith("python-httpx"):
                 self.send_response(418)
@@ -55,7 +62,7 @@ def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[s
 
 
 def test_public_ocr_executes_the_compiled_extension_without_python_fallback(
-    native_only_ocr_server: tuple[ThreadingHTTPServer, list[dict[str, object]]],
+    native_only_ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]],
 ) -> None:
     server, requests = native_only_ocr_server
     address: Final = server.server_address
@@ -71,10 +78,8 @@ def test_public_ocr_executes_the_compiled_extension_without_python_fallback(
 
     assert response.pages[0].markdown == "native OCR response"
     assert len(requests) == 1
-    headers: Final = requests[0]["headers"]
-    assert isinstance(headers, dict)
-    assert not str(headers.get("user-agent", "")).startswith("python-httpx")
-    assert requests[0]["body"] == {
+    assert not requests[0].headers.get("user-agent", "").startswith("python-httpx")
+    assert requests[0].body == {
         "model": "mistral-ocr-latest",
         "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
     }

tests/test_litellm_rust/test_provenance.py

@@ -1,20 +0,0 @@
-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

@@ -1,115 +0,0 @@
-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