test(rust-bridge): pin Messages and Responses callback contracts on the native path (#46443)

* test(rust-bridge): pin Messages and Responses callback contracts on the native path

Each test runs the native route and asserts what the Python path does today, so a
known gap is a strict xfail that names its reason and the exception it fails with
Messages and Responses share one Surface and hook recorder in support/callbacks.py,
and each surface folder lists its own expected counts and gaps.

Body-level hooks from the function body (pre-request hook, websearch short-circuit,
Messages interceptors, planned and legacy agentic loops, streaming deployment hook)
are xfail until native runs them

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(rust-bridge): give every route folder the same one-concern-per-file layout

Route data, the route Surface and the server fixture move into each route folder, and support/ keeps only
route-neutral helpers. test_inference.py splits into chat_completions/ and responses/, ocr/test_requests.py
splits into requests, failures and settings, and Messages and Responses get streaming and body hook files.
Messages request tests that duplicate inference-messages request.rs are removed, and drop_params read
from litellm.drop_params is now covered

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(rust-bridge): expect the required-route decline error for unsupported Responses requests

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): import route data from the route folders in the secrets unit tests

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): import NoReturn in the OCR dispatch tests

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): move route-neutral helpers out of support/callbacks.py

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): cover deployment rejection and concurrent callbacks

* test(rust-bridge): assert native serving via the x-litellm-rust marker in the route folder tests

Co-Authored-By: Yujong Lee <yujong@berri.ai>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/test_litellm_rust/AGENTS.md

@@ -1 +1,19 @@
 This directory holds only the tests that cannot be written in the Rust code
+
+Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses the same file names, one concern per file
+
+- `test_dispatch.py` which path served the call
+- `test_requests.py` Python inputs that reach the wire
+- `test_failures.py` errors raised to the caller and rejection before provider IO
+- `test_settings.py` globals, env, named credentials and their precedence
+- `test_streaming.py` stream relay, early close and stream-time logging
+- `test_callbacks.py` the logger hook contract
+- `test_body_hooks.py` handler body hooks such as interceptors and agentic loops
+- `test_lifecycle.py` cancellation, concurrency, context and nested calls
+
+A concern that exists for only one route gets its own file there, never a provider-named one
+
+- Route data and the route's `call` live in `<route>/data.py`, and `<route>/conftest.py` owns the `server` fixture
+- `support/` holds only route-neutral helpers. Never parametrize one test across routes, except in feature folders such as `cache/`
+- Call the public API inside `support.routing.native_route(...)`. Call `_native` directly only for a case the public API cannot reach
+- Check the native path with `support.routing.assert_served_natively`

tests/test_litellm_rust/cache/support.py

@@ -16,9 +16,9 @@ from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
 from litellm.rust_bridge.public_call import NativeCall
 from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES
 from litellm.types.utils import ModelResponse
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
-from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE
 
 
 def payload(value: object) -> object:

tests/test_litellm_rust/cache/test_python_cache.py

@@ -21,14 +21,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
 from litellm.rust_bridge import runtime
 from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
 from litellm.rust_bridge.configuration import Rollout
-from litellm.rust_bridge.dispatch import call_hook
 from litellm.rust_bridge.public_call import NativeCall
 from litellm.types.caching import CachingSupportedCallTypes
-from tests.test_litellm_rust.support.cache import cache_key, collect, invoke, payload
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
+from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
-from tests.test_litellm_rust.test_inference import RESPONSES_RESPONSE
 
 pytestmark = pytest.mark.requires_rust_extension
 
@@ -350,8 +349,6 @@ async def test_cache_lookup_uses_backend_request_callback_semantics(
     route: Literal["chat", "messages", "responses"],
     legacy: bool,
 ) -> None:
-    from tests.test_litellm_rust.support.requests import request_body
-
     class Rewrite(RecordingLogger):
         temperature = 0.1
 

tests/test_litellm_rust/cache/v2/test_cache.py

@@ -13,11 +13,11 @@ from litellm.caching.caching_handler import (
 )
 from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.types.utils import ModelResponse
-from tests.test_litellm_rust.support.cache import cache_key, invoke, payload
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT
+from tests.test_litellm_rust.cache.support import cache_key, invoke, payload
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
-from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE
 
 pytestmark = pytest.mark.requires_rust_extension
 
@@ -153,7 +153,7 @@ async def test_v2_cache_does_not_force_native_responses_streaming(
     assert isinstance(response, AsyncIterator)
     chunks: Final = [chunk async for chunk in response]
     assert chunks[-1].type == "response.completed"
-    assert chunks[-1].response.output[0].content[0].text == "native response"
+    assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT
 
 
 @pytest.mark.asyncio

tests/test_litellm_rust/chat_completions/conftest.py

@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    return recording_server

tests/test_litellm_rust/chat_completions/data.py

@@ -0,0 +1,28 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+
+def _error(kind: str) -> dict[str, object]:
+    return {"type": "error", "error": {"type": kind, "message": kind}}
+
+
+REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32}
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_error")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+    return await call_natively(
+        Route.CHAT_COMPLETIONS,
+        litellm.completion,
+        litellm.acompletion,
+        REQUEST,
+        server,
+        asynchronous=asynchronous,
+        **kwargs,
+    )

tests/test_litellm_rust/chat_completions/test_callbacks.py

@@ -0,0 +1,65 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import CallTypes, ModelResponse
+from tests.test_litellm_rust.chat_completions.data import call
+from tests.test_litellm_rust.messages.data import MESSAGES_TEXT
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, request_body
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_returns_a_model_response_and_logs_success_once(server: RecordingServer, asynchronous: bool) -> None:
+    recorder: Final = RecordingLogger()
+
+    result: Final = await call(server, asynchronous=asynchronous, callbacks=[recorder], temperature=0.25)
+
+    assert_served_natively(server)
+    sent: Final = server.requests[0]
+    body: Final = _OBJECT.validate_python(sent.body)
+    assert sent.path == "/v1/messages"
+    assert body["temperature"] == 0.25
+    assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body
+    assert isinstance(result, ModelResponse)
+    assert result.choices[0].message.content == MESSAGES_TEXT
+    success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event")
+    assert len(success) == 1
+    assert success[0].response is result
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+            request_body(kwargs)["temperature"] = 0.75
+
+    await call(server, callbacks=[Edit()])
+
+    assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.75
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
+    class DropKeywords(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: CallTypes | None
+        ) -> dict[str, object]:
+            return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")}
+
+    litellm.callbacks.append(DropKeywords())
+    await call(server, temperature=0.25, extra_headers={"x-deleted": "remove this header"})
+
+    assert_served_natively(server)
+    sent: Final = server.requests[0]
+    assert "temperature" not in _OBJECT.validate_python(sent.body)
+    assert sent.headers.get("x-deleted") is None

tests/test_litellm_rust/chat_completions/test_failures.py

@@ -0,0 +1,69 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm import RateLimitError
+from tests.test_litellm_rust.chat_completions.data import RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+def failure_signal(asynchronous: bool) -> str:
+    return "async_log_failure_event" if asynchronous else "log_failure_event"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+    monkeypatch.setattr(litellm, "max_budget", 1.0)
+    monkeypatch.setattr(litellm, "_current_cost", 2.0)
+    monkeypatch.setattr(litellm, "ssl_verify", object())
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+
+    with pytest.raises(litellm.BudgetExceededError) as caught:
+        await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+    failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not server.requests
+    assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_provider_failure_is_terminal_and_shared_with_callbacks(
+    server: RecordingServer, asynchronous: bool
+) -> None:
+    recorder: Final = RecordingLogger()
+    server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429))
+
+    with pytest.raises(RateLimitError) as caught:
+        await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+    assert caught.value.status_code == 429
+    assert len(server.requests) == 1
+    failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_empty_messages_fail_before_provider_io(server: RecordingServer, asynchronous: bool) -> None:
+    server.expected_requests = 0
+
+    with pytest.raises(Exception, match="chat completions requires at least one message"):
+        await call(server, asynchronous=asynchronous, messages=[])
+
+    assert not server.requests

tests/test_litellm_rust/chat_completions/test_lifecycle.py

@@ -0,0 +1,37 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge import _native
+from litellm.rust_bridge.public_call import NativeCall
+from tests.test_litellm_rust.chat_completions.data import REQUEST
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+async def test_unstarted_call_has_no_provider_or_callback_effects(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+    monkeypatch.setattr(litellm, "ssl_verify", object())
+    monkeypatch.setattr(litellm, "max_budget", 1.0)
+    monkeypatch.setattr(litellm, "_current_cost", 2.0)
+    kwargs: Final = {
+        **REQUEST,
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        "callbacks": [recorder],
+    }
+
+    pending: Final = _native.acompletion(NativeCall(args=(), kwargs=kwargs, base={}))
+    assert asyncio.iscoroutine(pending)
+    pending.close()
+
+    assert not server.requests
+    assert not recorder.events

tests/test_litellm_rust/chat_completions/test_requests.py

@@ -0,0 +1,42 @@
+import asyncio
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from litellm.types.utils import ModelResponse
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_TEXT
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None:
+    temperature: Final = 0.35
+    arguments: Final = (
+        (MESSAGES_MODEL, list(MESSAGES), None, None, 12.0, temperature)
+        if asynchronous
+        else (MESSAGES_MODEL, list(MESSAGES), 12.0, temperature)
+    )
+    supplied: Final = {"base_url": server.base_url, "api_key": "test-key", "max_tokens": 32}
+
+    with native_route(Route.CHAT_COMPLETIONS):
+        result: Final = (
+            await litellm.acompletion(*arguments, **supplied)
+            if asynchronous
+            else await asyncio.to_thread(litellm.completion, *arguments, **supplied)
+        )
+
+    assert isinstance(result, ModelResponse)
+    assert result.choices[0].message.content == MESSAGES_TEXT
+    assert_served_natively(server)
+    body: Final = _OBJECT.validate_python(server.requests[0].body)
+    assert body["temperature"] == temperature
+    assert body["max_tokens"] == supplied["max_tokens"]

tests/test_litellm_rust/chat_completions/test_settings.py

@@ -0,0 +1,82 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.models.credentials import CredentialItem
+from litellm.types.utils import CallTypes, ModelResponse
+from tests.test_litellm_rust.chat_completions.data import call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("from_credentials", (False, True))
+async def test_resource_setup_uses_deployment_hook_arguments(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool
+) -> None:
+    invalid_settings: Final = {"ssl_verify": object()}
+    credential: Final = CredentialItem(
+        credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
+    )
+    monkeypatch.setattr(litellm, "credential_list", [credential])
+
+    class Prepare(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: CallTypes | None
+        ) -> dict[str, object]:
+            return {
+                **kwargs,
+                **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
+            }
+
+    litellm.callbacks.append(Prepare())
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+
+    with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
+        await call(server, callbacks=[recorder])
+
+    failure: Final = await recorder.wait_for_async("async_log_failure_event")
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not server.requests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
+async def test_connection_settings_reach_the_provider(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str
+) -> None:
+    key: Final = "selected-key"
+    explicit: Final = source in ("explicit", "base_url")
+    monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
+    monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None))
+    monkeypatch.setattr(
+        litellm,
+        "api_base",
+        None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url),
+    )
+    monkeypatch.setenv("ANTHROPIC_API_KEY", key if source == "environment" else "unused")
+    for name in ("ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"):
+        monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1")
+
+    result: Final = await call(
+        server,
+        asynchronous=asynchronous,
+        api_key=key if explicit else ("" if source == "empty" else None),
+        api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None),
+        **({"base_url": server.base_url} if source == "base_url" else {}),
+    )
+
+    assert isinstance(result, ModelResponse)
+    assert_served_natively(server)
+    assert server.requests[0].headers["x-api-key"] == key

tests/test_litellm_rust/conftest.py

@@ -23,12 +23,12 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, re
 
 
 @pytest_asyncio.fixture(autouse=True, loop_scope="function")
-async def isolate_ocr_test_state() -> AsyncIterator[None]:
+async def isolate_test_state() -> AsyncIterator[None]:
     with ExitStack() as stack:
         stack.enter_context(isolated_callback_registries())
         stack.enter_context(rebound(litellm, "cache", None))  # test-quality-ok: isolate process-global cache
         stack.enter_context(rebound(_CONFIGURATION, "override", None))
-        executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging")
+        executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging")
         stack.enter_context(rebound(litellm_logging, "executor", executor))
         stack.enter_context(rebound(utils, "executor", executor))
         stack.enter_context(rebound(thread_pool_executor, "executor", executor))

tests/test_litellm_rust/messages/conftest.py

@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+    return recording_server

tests/test_litellm_rust/messages/data.py

@@ -0,0 +1,64 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5"
+MESSAGES: Final = ({"role": "user", "content": "Hello"},)
+MESSAGES_TEXT: Final = "Hello from native Messages"
+MESSAGES_RESPONSE: Final = {
+    "id": "msg_native",
+    "type": "message",
+    "role": "assistant",
+    "model": "claude-sonnet-5",
+    "content": [{"type": "text", "text": MESSAGES_TEXT}],
+    "stop_reason": "end_turn",
+    "stop_sequence": None,
+    "usage": {"input_tokens": 5, "output_tokens": 4},
+}
+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": MESSAGES_TEXT}},
+    ),
+    ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+    (
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": 4},
+        },
+    ),
+    ("message_stop", {"type": "message_stop"}),
+)
+
+REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 64}
+
+
+def _error(kind: str) -> dict[str, object]:
+    return {"type": "error", "error": {"type": kind, "message": kind}}
+
+
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_error")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+    return await call_natively(
+        Route.MESSAGES,
+        litellm.anthropic.messages.create,
+        litellm.anthropic.messages.acreate,
+        REQUEST,
+        server,
+        asynchronous=asynchronous,
+        **kwargs,
+    )
+
+
+BODY_HOOKS: Final = gap("native Messages does not run the handler body hooks")

tests/test_litellm_rust/messages/test_body_hooks.py

@@ -0,0 +1,137 @@
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.anthropic.pass_through.messages import interceptors
+from litellm.llms.anthropic.pass_through.messages.interceptors.base import MessagesInterceptor
+from tests.test_litellm_rust.messages.data import BODY_HOOKS, MESSAGES_RESPONSE, REQUEST, call
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id
+from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_pre_request_hook_rewrites_the_provider_request(server: RecordingServer) -> None:
+    class Rewrite(CustomLogger):
+        async def async_pre_request_hook(
+            self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object]
+        ) -> dict[str, object]:
+            return {**kwargs, "system": "body hook instruction"}
+
+    with rebound(litellm, "callbacks", [Rewrite()]):
+        await call(server)
+
+    assert_served_natively(server)
+    assert isinstance(server.requests[0].body, dict)
+    assert server.requests[0].body.get("system") == "body hook instruction"
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_websearch_short_circuit_answers_without_the_provider(server: RecordingServer) -> None:
+    class ShortCircuit(WebSearchInterceptionLogger):
+        async def try_short_circuit_search(
+            self,
+            model: str,
+            messages: list[dict[str, object]],
+            tools: list[dict[str, object]] | None,
+            custom_llm_provider: str | None,
+            kwargs: dict[str, object] | None = None,
+        ) -> object:
+            return {**MESSAGES_RESPONSE, "id": "short-circuited"}
+
+    with rebound(litellm, "callbacks", [ShortCircuit()]):
+        response: Final = await call(server)
+
+    assert response_id(response) == "short-circuited"
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_interceptor_replaces_the_provider_call(server: RecordingServer) -> None:
+    class Replace(MessagesInterceptor):
+        def can_handle(self, tools: list[dict[str, object]] | None, custom_llm_provider: str | None) -> bool:
+            return True
+
+        async def handle(self, **kwargs: object) -> object:
+            return {**MESSAGES_RESPONSE, "id": "intercepted"}
+
+    with rebound(interceptors, "_interceptors", [Replace()]):
+        response: Final = await call(server)
+
+    assert response_id(response) == "intercepted"
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None:
+    loop: Final = PlannedAgenticLoop()
+    server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE))
+    server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "id": "followup"})
+
+    with rebound(litellm, "callbacks", [loop]):
+        response: Final = await call(server)
+        await drain_logging()
+
+    assert_served_natively(server, count=2)
+    assert isinstance(server.requests[1].body, dict)
+    assert server.requests[1].body["messages"] == [{"role": "user", "content": "followup"}]
+    assert response_id(response) == "agentic-replaced"
+    assert agentic_hooks(loop) == (
+        "async_should_run_agentic_loop",
+        "async_build_agentic_loop_plan",
+        "async_should_run_agentic_loop",
+        "async_post_agentic_loop_response_hook",
+    )
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None:
+    loop: Final = LegacyAgenticLoop()
+
+    with rebound(litellm, "callbacks", [loop]):
+        response: Final = await call(server)
+
+    assert_served_natively(server)
+    assert response_id(response) == "agentic-replaced"
+    assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop")
+
+
+@pytest.mark.asyncio
+async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None:
+    observed: Final[list[object]] = []
+    logging_obj: Final = Logging(
+        model=REQUEST["model"],
+        messages=[],
+        stream=False,
+        start_time=datetime.now(),
+        call_type="anthropic_messages",
+        litellm_call_id=CALL_ID,
+        function_id=CALL_ID,
+    )
+    logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={})
+
+    await call(server, litellm_logging_obj=logging_obj)
+
+    assert_served_natively(server)
+    assert observed == []

tests/test_litellm_rust/messages/test_callbacks.py

@@ -1,197 +1,138 @@
-from collections.abc import AsyncIterator, Iterator
+import asyncio
+from collections.abc import Iterator
+from contextvars import ContextVar
+from datetime import datetime
 from typing import Final
 
 import pytest
+from pydantic import TypeAdapter
 
 import litellm
+from litellm.integrations.argilla import ArgillaLogger
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.rust_bridge import catalog
-from litellm.rust_bridge.catalog import Route, RouteRule
-from litellm.rust_bridge.configuration import Rollout
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from litellm.types.integrations.argilla import ArgillaItem
+from litellm.types.utils import ModelResponse, StandardLoggingPayload
+from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
+from tests.test_litellm_rust.support.gaps import gap
 from tests.test_litellm_rust.support.isolation import rebound
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import (
-    MESSAGES,
-    MESSAGES_EVENTS,
-    MESSAGES_MODEL,
-    MESSAGES_RESPONSE,
-    request_body,
-)
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id, with_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
 
 pytestmark = pytest.mark.requires_rust_extension
 
-STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
-
-
-@pytest.fixture(autouse=True)
-def opt_messages_into_rust() -> Iterator[None]:
-    with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)):
-        yield
+SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers")
+_OBJECT: Final = TypeAdapter(dict[str, object])
+_HEADERS: Final = TypeAdapter(dict[str, str])
 
 
 @pytest.fixture
-def messages_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
-    return recording_server
-
-
-def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
-    return {
-        "model": MESSAGES_MODEL,
-        "messages": [dict(message) for message in MESSAGES],
-        "max_tokens": 64,
-        "api_key": "test-key",
-        "api_base": server.base_url,
-        **kwargs,
-    }
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
 
 
 @pytest.mark.asyncio
-async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response(
-    messages_server: RecordingServer,
+@pytest.mark.parametrize(
+    "hook,expected",
+    (
+        pytest.param("log_pre_api_call", 1),
+        pytest.param("log_post_api_call", 1),
+        pytest.param("async_log_success_event", 1),
+        pytest.param("async_logging_hook", 1),
+        pytest.param("redact_standard_logging_payload_from_model_call_details", 1, id="redact"),
+        pytest.param("async_pre_call_deployment_hook", 1),
+        pytest.param("async_post_call_success_deployment_hook", 1),
+        pytest.param("log_success_event", 0),
+        pytest.param("logging_hook", 0),
+    ),
+)
+async def test_success_hook_fires_the_expected_number_of_times(
+    server: RecordingServer, hook: str, expected: int
 ) -> None:
+    recorder: Final = HookRecorder()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        response: Final = await call(server)
+        await recorder.wait_for_async("async_log_success_event")
+
+    assert_served_natively(server, response)
+    assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "hook,expected",
+    (
+        pytest.param("async_log_failure_event", 1),
+        pytest.param("async_post_call_failure_deployment_hook", 1),
+        pytest.param("log_model_group_rate_limit_error", 1),
+        pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+    ),
+)
+async def test_failure_hook_fires_the_expected_number_of_times(
+    server: RecordingServer, hook: str, expected: int
+) -> None:
+    server.default_response = ResponseSpec(body=RATE_LIMITED, status=429)
+    recorder: Final = HookRecorder()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        with pytest.raises(litellm.RateLimitError):
+            await call(server, metadata={"model_group_size": 1})
+        await recorder.wait_for_async("async_log_failure_event")
+
+    assert_served_natively(server)
+    assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None:
+    server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
     recorder: Final = RecordingLogger()
 
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success")
-    )
+    with rebound(litellm, "callbacks", [recorder]):
+        with pytest.raises(litellm.BadRequestError) as raised:
+            await call(server)
+        events: Final = await recorder.wait_for_async("async_log_failure_event")
 
-    assert_served_natively(messages_server, response)
-    assert response["content"] == MESSAGES_RESPONSE["content"]
-    sent: Final = messages_server.requests[0]
-    assert sent.path == "/v1/messages"
-    assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False}
-    pre_call: Final = recorder.wait_for("log_pre_api_call")
-    assert request_body(pre_call[0].kwargs) == sent.body
-    success: Final = await recorder.wait_for_async("async_log_success_event")
-    assert len(success) == 1
-    assert success[0].call_type == "anthropic_messages"
-    assert success[0].kwargs["litellm_call_id"] == "messages-success"
-    assert success[0].response.choices[0].message.content == "Hello from native Messages"
+    assert_served_natively(server)
+    assert isinstance(events[0].kwargs, dict)
+    assert events[0].kwargs["exception"] is raised.value
 
 
 @pytest.mark.asyncio
-async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None:
+async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    response: Final = await call(server, callbacks=[recorder])
+
+    assert_served_natively(server, response)
+    pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
+    assert isinstance(pre_call, dict)
+    assert request_body(pre_call) == server.requests[0].body
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
     class Edit(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
+        def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
             request_body(kwargs)["temperature"] = 0.25
 
-    await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()]))
+    response: Final = await call(server, callbacks=[Edit()])
 
-    assert messages_server.requests[0].body["temperature"] == 0.25
+    assert_served_natively(server, response)
+    assert isinstance(server.requests[0].body, dict)
+    assert server.requests[0].body["temperature"] == 0.25
 
 
 @pytest.mark.asyncio
-async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error(
-    messages_server: RecordingServer,
-) -> None:
-    messages_server.enqueue(
-        ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400)
-    )
-    observed: Final = []
-
-    class Observe(CustomLogger):
-        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            observed.append(("sync", kwargs["exception"]))
-
-        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
-            observed.append(("async", kwargs["exception"]))
-
-    with pytest.raises(litellm.BadRequestError) as raised:
-        await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()]))
-
-    assert len(messages_server.requests) == 1
-    assert [phase for phase, _ in observed] == ["sync", "async"]
-    assert all(error is raised.value for _, error in observed)
-
-
-def sse_payload() -> bytes:
-    return b"".join(STREAM.payloads())
-
-
-@pytest.mark.asyncio
-async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
-    messages_server: RecordingServer,
-) -> None:
-    messages_server.enqueue(STREAM)
+async def test_pre_api_call_sees_the_shaped_optional_params(server: RecordingServer) -> None:
     recorder: Final = RecordingLogger()
 
-    stream: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, stream=True, callbacks=[recorder])
-    )
-    assert isinstance(stream, AsyncIterator)
-    first: Final = await anext(stream)
-    await drain_logging()
-    assert "async_log_success_event" not in recorder.names
-    rest: Final = [chunk async for chunk in stream]
+    await call(server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True)
 
-    assert first + b"".join(rest) == sse_payload()
-    assert_served_natively(messages_server, stream)
-    assert messages_server.requests[0].body["stream"] is True
-    success: Final = await recorder.wait_for_async("async_log_success_event")
-    assert len(success) == 1
-    assert success[0].kwargs["stream"] is True
-    assert success[0].kwargs["completion_start_time"] is not None
-    assert "log_failure_event" not in recorder.names
-
-
-@pytest.mark.asyncio
-async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered(
-    messages_server: RecordingServer,
-) -> None:
-    messages_server.enqueue(STREAM)
-    recorder: Final = RecordingLogger()
-
-    stream: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, stream=True, callbacks=[recorder])
-    )
-    assert isinstance(stream, AsyncIterator)
-    await anext(stream)
-    await stream.aclose()
-
-    success: Final = await recorder.wait_for_async("async_log_success_event")
-    assert len(success) == 1
-    with pytest.raises(StopAsyncIteration):
-        await anext(stream)
-
-
-def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once(
-    messages_server: RecordingServer,
-) -> None:
-    messages_server.enqueue(STREAM)
-    recorder: Final = RecordingLogger()
-
-    stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder]))
-    assert isinstance(stream, Iterator)
-
-    assert b"".join(stream) == sse_payload()
-    assert_served_natively(messages_server, stream)
-    assert len(recorder.wait_for("async_log_success_event")) == 1
-
-
-def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None:
-    recorder: Final = RecordingLogger()
-
-    response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder]))
-
-    assert_served_natively(messages_server, response)
-    assert response["content"] == MESSAGES_RESPONSE["content"]
-    assert len(recorder.wait_for("log_success_event")) == 1
-
-
-@pytest.mark.asyncio
-async def test_native_messages_pre_call_sees_the_shaped_optional_params(
-    messages_server: RecordingServer,
-) -> None:
-    recorder: Final = RecordingLogger()
-
-    await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True)
-    )
-
-    sent: Final = messages_server.requests[0].body
+    sent: Final = server.requests[0].body
+    assert isinstance(sent, dict)
     assert not {"temperature", "top_k"} & sent.keys()
     pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
     assert isinstance(pre_call, dict)
@@ -202,65 +143,287 @@ async def test_native_messages_pre_call_sees_the_shaped_optional_params(
 
 
 @pytest.mark.asyncio
-async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None:
-    class Broken(CustomLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            raise RuntimeError("logger exploded")
+async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None:
+    class Rewrite(HookRecorder):
+        async def async_logging_hook(
+            self, kwargs: dict[str, object], result: object, call_type: str
+        ) -> tuple[dict[str, object], object]:
+            return {**kwargs, "logging_hook_ran": True}, result
 
-    response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()]))
+        def redact_standard_logging_payload_from_model_call_details(
+            self, model_call_details: dict[str, object]
+        ) -> dict[str, object]:
+            return {**model_call_details, "redacted": True}
 
-    assert_served_natively(messages_server, response)
-    assert response["content"] == MESSAGES_RESPONSE["content"]
+    recorder: Final = Rewrite()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        await call(server)
+        success: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert isinstance(success[0].kwargs, dict)
+    assert success[0].kwargs["logging_hook_ran"] is True
+    assert success[0].kwargs["redacted"] is True
 
 
 @pytest.mark.asyncio
-async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events(
-    messages_server: RecordingServer,
+async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order(
+    server: RecordingServer,
 ) -> None:
-    messages_server.enqueue(STREAM)
+    class Rewrite(CustomLogger):
+        def __init__(self, marker: str) -> None:
+            super().__init__()
+            self.marker = marker
+
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            headers: Final = _HEADERS.validate_python(kwargs["extra_headers"])
+            return {
+                **kwargs,
+                "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker},
+            }
+
+    with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]):
+        response: Final = await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"})
+
+    assert_served_natively(server, response)
+    assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second"
+    assert server.requests[0].headers["x-kept"] == "unchanged"
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None:
+    rejection: Final = litellm.BadRequestError(
+        message="deployment rejected", model="test-model", llm_provider="anthropic"
+    )
     recorder: Final = RecordingLogger()
 
-    stream: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, stream=True, callbacks=[recorder])
-    )
-    assert isinstance(stream, AsyncIterator)
-    async for _ in stream:
-        pass
+    class Reject(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            raise rejection
 
-    success: Final = await recorder.wait_for_async("async_log_success_event")
-    usage: Final = success[0].response.usage
-    assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"]
-    assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
-    assert success[0].response.choices[0].message.content == "Hello from native Messages"
+    with rebound(litellm, "callbacks", [Reject(), recorder]):
+        with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised:
+            await call(server)
+        failures: Final = await recorder.wait_for_async("async_log_failure_event")
+
+    assert raised.value is rejection
+    assert server.requests == []
+    assert len(failures) == 1
+    assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value
+    assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names)
 
 
 @pytest.mark.asyncio
-async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_hook(
-    messages_server: RecordingServer,
-) -> None:
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
     class DropKeywords(CustomLogger):
         async def async_pre_call_deployment_hook(
             self, kwargs: dict[str, object], call_type: object
         ) -> dict[str, object]:
             return {name: value for name, value in kwargs.items() if name not in ("system", "extra_headers")}
 
-    call_arguments: Final = arguments(
-        messages_server,
-        system="remove this instruction",
-        extra_headers={"x-deleted": "remove this header"},
-        headers={"x-kept": "keep this header"},
-        stream=False,
-    )
     with rebound(litellm, "callbacks", [DropKeywords()]):
-        response: Final = await litellm.anthropic.messages.acreate(**call_arguments)
+        response: Final = await call(
+            server,
+            system="remove this instruction",
+            extra_headers={"x-deleted": "remove this header"},
+            headers={"x-kept": "keep this header"},
+        )
 
-    assert_served_natively(messages_server, response)
-    sent: Final = messages_server.requests[0]
-    assert sent.body == {
-        "model": MESSAGES_MODEL.removeprefix("anthropic/"),
-        "messages": call_arguments["messages"],
-        "max_tokens": call_arguments["max_tokens"],
-        "stream": call_arguments["stream"],
-    }
-    assert sent.headers.get("x-deleted") is None
-    assert sent.headers.get("x-kept") == "keep this header"
+    assert_served_natively(server, response)
+    sent: Final = server.requests[0]
+    assert isinstance(sent.body, dict)
+    assert "system" not in sent.body
+    assert "x-deleted" not in sent.headers
+    assert sent.headers["x-kept"] == "keep this header"
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    class Replace(CustomLogger):
+        async def async_post_call_success_deployment_hook(
+            self, request_data: dict[str, object], response: object, call_type: object
+        ) -> object:
+            return with_id(response, "replaced")
+
+    with rebound(litellm, "callbacks", [Replace(), recorder]):
+        response: Final = await call(server)
+        success: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert_served_natively(server, response)
+    assert response_id(response) == "replaced"
+    assert len(success) == 1
+    assert isinstance(success[0].response, ModelResponse)
+    assert success[0].response.id == response_id(response)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None:
+    context: Final = ContextVar("messages-callback-context", default="outside")
+    first_entered: Final = asyncio.Event()
+    second_entered: Final = asyncio.Event()
+    server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_first"}))
+    server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_second"}))
+
+    class Observe(RecordingLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            arrived, other = (
+                (first_entered, second_entered)
+                if kwargs["litellm_call_id"] == "first"
+                else (second_entered, first_entered)
+            )
+            arrived.set()
+            await asyncio.wait_for(other.wait(), timeout=5)
+            return kwargs
+
+        def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+            self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()})
+
+        async def async_log_success_event(
+            self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+        ) -> None:
+            self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj)
+
+    recorder: Final = Observe()
+
+    async def invoke(marker: str) -> object:
+        token: Final = context.set(marker)
+        try:
+            return await call(server, litellm_call_id=marker, metadata={"marker": marker})
+        finally:
+            context.reset(token)
+
+    with rebound(litellm, "callbacks", [recorder]):
+        responses: Final = await asyncio.gather(invoke("first"), invoke("second"))
+        success: Final = await recorder.wait_for_async("async_log_success_event", count=2)
+
+    assert_served_natively(server, *responses)
+    expected: Final = {"first": responses[0], "second": responses[1]}
+    pre_call: Final = recorder.wait_for("log_pre_api_call", count=2)
+    assert len(pre_call) == len(success) == 2
+    assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys()
+    for event in (*pre_call, *success):
+        details: Final = _OBJECT.validate_python(event.kwargs)
+        call_id: Final = str(details["litellm_call_id"])
+        params: Final = _OBJECT.validate_python(details["litellm_params"])
+        assert details["caller_context"] == call_id
+        assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id
+    for event in success:
+        logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"])
+        assert isinstance(event.response, ModelResponse)
+        assert event.response.id == response_id(expected[logged_call_id])
+    assert context.get() == "outside"
+
+
+@pytest.mark.asyncio
+async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None:
+    class Broken(CustomLogger):
+        def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+            raise RuntimeError("logger exploded")
+
+    response: Final = await call(server, callbacks=[Broken()])
+
+    assert_served_natively(server, response)
+    assert response_id(response) == MESSAGES_RESPONSE["id"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "event,asynchronous,dynamic,expected",
+    (
+        pytest.param("input", False, False, 1, id="global-input"),
+        pytest.param("success", True, False, 1, id="global-async-success"),
+        pytest.param("success", True, True, 1, id="request-async-success"),
+        pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+        pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+        pytest.param("failure", True, False, 1, id="global-async-failure"),
+        pytest.param("failure", False, False, 1, id="global-sync-failure"),
+        pytest.param("failure", False, True, 1, id="request-sync-failure"),
+    ),
+)
+async def test_function_callback_fires_the_expected_number_of_times(
+    server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int
+) -> None:
+    recorder: Final = RecordingLogger()
+
+    def callback(kwargs: dict[str, object], *_: object) -> None:
+        recorder._record(event, kwargs)
+
+    async def async_callback(kwargs: dict[str, object], *_: object) -> None:
+        recorder._record(event, kwargs)
+
+    registered: Final = async_callback if asynchronous else callback
+    global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback"
+    request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {}
+    if event == "failure":
+        server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+
+    with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]):
+        try:
+            await call(server, **request_callbacks)
+        except litellm.BadRequestError:
+            assert event == "failure"
+        await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event")
+
+    assert_served_natively(server)
+    assert recorder.names.count(event) == expected
+
+
+@pytest.mark.asyncio
+async def test_pre_call_rule_rejects_before_provider_io(server: RecordingServer) -> None:
+    checked: Final[list[str]] = []
+
+    def reject(text: str) -> bool:
+        checked.append(text)
+        return False
+
+    with rebound(litellm, "pre_call_rules", [reject]):
+        with pytest.raises(litellm.APIResponseValidationError):
+            await call(server)
+
+    assert checked == ["Hello"]
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None:
+    datasets: Final = RecordingLogger()
+
+    class Drop(CustomLogger):
+        async def async_dataset_hook(
+            self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None
+        ) -> None:
+            datasets._record("async_dataset_hook", response=logged_data)
+
+    server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]}))
+    with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}):
+        sink: Final = ArgillaLogger(
+            argilla_api_key="test-key",
+            argilla_base_url=server.base_url,
+            argilla_dataset_name="callback-dataset",
+            flush_interval=3600,
+        )
+        with rebound(litellm, "callbacks", [sink, Drop()]):
+            await call(server)
+            await datasets.wait_for_async("async_dataset_hook")
+            await sink.flush_queue()
+
+    assert [request.method for request in server.requests] == ["GET", "POST"]
+
+
+@pytest.mark.asyncio
+async def test_sync_call_returns_the_provider_message(server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    response: Final = await call(server, asynchronous=False, callbacks=[recorder])
+
+    assert_served_natively(server, response)
+    assert response_id(response) == MESSAGES_RESPONSE["id"]
+    assert len(await recorder.wait_for_async("log_success_event")) == 1

tests/test_litellm_rust/messages/test_request_shaping.py

@@ -1,311 +0,0 @@
-"""The native Messages route shapes the wire request the way the Python handler does.
-
-Model capability expectations come from model_prices_and_context_window.json (Claude Sonnet 5 is an
-adaptive-thinking model without sampling params; Claude Haiku 4.5 is a legacy-thinking model), read at
-2026-09-24; the cost map is LiteLLM's own file.
-"""
-
-from collections.abc import Iterator
-from typing import Final
-
-import pytest
-
-import litellm
-from litellm.rust_bridge import catalog
-from litellm.rust_bridge.catalog import Route, RouteRule
-from litellm.rust_bridge.configuration import Rollout
-from tests.test_litellm_rust.support.isolation import rebound
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_RESPONSE
-from tests.test_litellm_rust.support.served import assert_served_natively
-
-pytestmark = pytest.mark.requires_rust_extension
-
-ADAPTIVE_MODEL: Final = "anthropic/claude-sonnet-5"
-LEGACY_THINKING_MODEL: Final = "anthropic/claude-haiku-4-5"
-
-
-@pytest.fixture(autouse=True)
-def opt_messages_into_rust() -> Iterator[None]:
-    with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)):
-        yield
-
-
-@pytest.fixture
-def messages_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
-    return recording_server
-
-
-def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
-    return {
-        "model": ADAPTIVE_MODEL,
-        "messages": [dict(message) for message in MESSAGES],
-        "max_tokens": 8192,
-        "api_key": "test-key",
-        "api_base": server.base_url,
-        **kwargs,
-    }
-
-
-def sent(server: RecordingServer, response: object) -> tuple[dict[str, object], dict[str, str]]:
-    assert_served_natively(server, response)
-    request: Final = server.requests[0]
-    assert isinstance(request.body, dict)
-    return request.body, request.headers
-
-
-@pytest.mark.asyncio
-async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire(
-    messages_server: RecordingServer,
-) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high"))
-
-    body, _ = sent(messages_server, response)
-    assert "reasoning_effort" not in body
-    assert body["thinking"] == {"type": "adaptive", "display": "summarized"}
-    assert body["output_config"] == {"effort": "high"}
-
-
-@pytest.mark.asyncio
-async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for_a_legacy_model(
-    messages_server: RecordingServer,
-) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(
-            messages_server,
-            model=LEGACY_THINKING_MODEL,
-            max_tokens=3000,
-            thinking={"type": "adaptive"},
-            output_config={"effort": "high"},
-            temperature=0,
-        )
-    )
-
-    body, _ = sent(messages_server, response)
-    assert body["thinking"] == {"type": "enabled", "budget_tokens": 2999}
-    assert "output_config" not in body
-    assert "temperature" not in body
-
-
-@pytest.mark.asyncio
-async def test_removed_sampling_params_are_dropped_under_drop_params(messages_server: RecordingServer) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, temperature=0.2, top_p=0.9, top_k=5, drop_params=True)
-    )
-
-    body, _ = sent(messages_server, response)
-    assert not {"temperature", "top_p", "top_k"} & body.keys()
-
-
-@pytest.mark.asyncio
-async def test_removed_sampling_params_are_rejected_without_drop_params(messages_server: RecordingServer) -> None:
-    messages_server.expected_requests = 0
-
-    with pytest.raises(litellm.BadRequestError, match="does not support top_k=5"):
-        await litellm.anthropic.messages.acreate(**arguments(messages_server, top_k=5))
-
-    assert messages_server.requests == []
-
-
-@pytest.mark.asyncio
-async def test_replayed_history_is_sanitized_before_it_reaches_the_provider(
-    messages_server: RecordingServer,
-) -> None:
-    history: Final = [
-        {"role": "user", "content": "run it"},
-        {
-            "role": "assistant",
-            "content": [
-                {"type": "text", "text": ""},
-                {
-                    "type": "tool_use",
-                    "id": "functions.Bash:0",
-                    "name": "Bash",
-                    "input": {},
-                    "provider_specific_fields": {"x": 1},
-                },
-            ],
-        },
-        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]},
-    ]
-
-    response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history))
-
-    body, _ = sent(messages_server, response)
-    assert body["messages"] == [
-        {"role": "user", "content": "run it"},
-        {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]},
-        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]},
-    ]
-
-
-@pytest.mark.asyncio
-async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_server: RecordingServer) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(
-            messages_server,
-            output_format={"type": "json_schema", "schema": {"type": "object"}},
-            extra_headers={"anthropic-beta": "web-search-2025-03-05"},
-        )
-    )
-
-    _, headers = sent(messages_server, response)
-    assert headers["anthropic-beta"] == "structured-outputs-2025-11-13,web-search-2025-03-05"
-
-
-@pytest.mark.asyncio
-async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(messages_server: RecordingServer) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, api_key="sk-ant-oat01-token")
-    )
-
-    _, headers = sent(messages_server, response)
-    assert "x-api-key" not in headers
-    assert headers["authorization"] == "Bearer sk-ant-oat01-token"
-    assert headers["anthropic-beta"] == "oauth-2025-04-20"
-    assert headers["anthropic-dangerous-direct-browser-access"] == "true"
-
-
-@pytest.mark.asyncio
-async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_server: RecordingServer) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, metadata={"user_id": "u-1", "trace_id": "internal"})
-    )
-
-    body, _ = sent(messages_server, response)
-    assert body["metadata"] == {"user_id": "u-1"}
-
-
-@pytest.mark.asyncio
-async def test_additional_drop_params_remove_nested_fields_from_the_wire(messages_server: RecordingServer) -> None:
-    tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
-
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(messages_server, tools=tools, additional_drop_params=["tools[*].input_examples"])
-    )
-
-    body, _ = sent(messages_server, response)
-    assert body["tools"] == [{"name": "lookup", "input_schema": {"type": "object"}}]
-
-
-@pytest.mark.asyncio
-async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(messages_server: RecordingServer) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(
-            messages_server,
-            provider_specific_header=[
-                {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}},
-                {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}},
-            ],
-        )
-    )
-
-    _, headers = sent(messages_server, response)
-    assert headers["x-scoped"] == "yes"
-    assert "x-other" not in headers
-
-
-@pytest.mark.asyncio
-async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers(
-    messages_server: RecordingServer,
-) -> None:
-    response: Final = await litellm.anthropic.messages.acreate(
-        **arguments(
-            messages_server,
-            headers={"x-priority": "forwarded", "x-forwarded-only": "kept"},
-            extra_headers={"x-priority": "extra", "x-extra-only": "kept"},
-            provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}},
-        )
-    )
-
-    _, headers = sent(messages_server, response)
-    assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == {
-        "x-priority": "scoped",
-        "x-forwarded-only": "kept",
-        "x-extra-only": "kept",
-    }
-
-
-@pytest.mark.asyncio
-async def test_non_string_metadata_user_id_is_rejected_before_the_provider_call(
-    messages_server: RecordingServer,
-) -> None:
-    messages_server.expected_requests = 0
-
-    with pytest.raises(litellm.BadRequestError, match=r"metadata\.user_id must be a string"):
-        await litellm.anthropic.messages.acreate(**arguments(messages_server, metadata={"user_id": 123}))
-
-    assert messages_server.requests == []
-
-
-@pytest.mark.asyncio
-async def test_native_messages_observes_runtime_capabilities_and_separate_caller_settings(
-    messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
-) -> None:
-    from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
-    from litellm.rust_bridge.public_call import NativeCall
-
-    native: Final = NATIVE_AMESSAGES.load()
-    assert native is not None
-    model: Final = "claude-test-runtime-capabilities"
-    messages_server.expected_requests = 2
-    request: Final = NativeCall(
-        args=(),
-        kwargs={"temperature": 0.2, "drop_params": True},
-        base={
-            "model": model,
-            "messages": MESSAGES,
-            "max_tokens": 16,
-            "stream": None,
-            "api_key": "test-key",
-            "api_base": messages_server.base_url,
-            "custom_llm_provider": "anthropic",
-            "temperature": 0.2,
-            "drop_params": True,
-        },
-    )
-    monkeypatch.setitem(
-        litellm.model_cost,
-        model,
-        {
-            "litellm_provider": "anthropic",
-            "mode": "chat",
-            "supports_sampling_params": True,
-        },
-    )
-    first: Final = await native(request)
-    monkeypatch.setitem(
-        litellm.model_cost,
-        model,
-        {
-            "litellm_provider": "anthropic",
-            "mode": "chat",
-            "supports_sampling_params": False,
-        },
-    )
-    second: Final = await native(request)
-
-    assert isinstance(first, dict)
-    assert isinstance(second, dict)
-    assert first["id"] == second["id"] == MESSAGES_RESPONSE["id"]
-    assert len(messages_server.requests) == 2
-    first_body: Final = messages_server.requests[0].body
-    second_body: Final = messages_server.requests[1].body
-    assert isinstance(first_body, dict)
-    assert isinstance(second_body, dict)
-    assert first_body["temperature"] == request.kwargs["temperature"]
-    assert second_body == {name: value for name, value in first_body.items() if name != "temperature"}
-
-
-@pytest.mark.asyncio
-async def test_native_messages_reads_optional_positional_body_parameters(messages_server: RecordingServer) -> None:
-    metadata: Final = {"user_id": "caller"}
-    args: Final = (16, MESSAGES, "anthropic/claude-test", metadata, None, False, "Be brief", 0.25)
-    kwargs: Final = {"api_key": "test-key", "api_base": messages_server.base_url}
-    response: Final = await litellm.anthropic.messages.acreate(*args, **kwargs)
-
-    body, _ = sent(messages_server, response)
-    assert body["temperature"] == args[7]
-    assert body["system"] == args[6]
-    assert body["metadata"] == metadata

tests/test_litellm_rust/messages/test_requests.py

@@ -0,0 +1,150 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_RESPONSE, call
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+CAPABILITY_MODEL: Final = "claude-test-runtime-capabilities"
+
+
+def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]:
+    assert_served_natively(server)
+    request: Final = server.requests[0]
+    assert isinstance(request.body, dict)
+    return request.body, request.headers
+
+
+def capabilities(monkeypatch: pytest.MonkeyPatch, *, sampling: bool) -> None:
+    monkeypatch.setitem(
+        litellm.model_cost,
+        CAPABILITY_MODEL,
+        {"litellm_provider": "anthropic", "mode": "chat", "supports_sampling_params": sampling},
+    )
+
+
+@pytest.mark.asyncio
+async def test_replayed_tool_ids_are_sanitized_before_they_reach_the_provider(server: RecordingServer) -> None:
+    history: Final = [
+        {"role": "user", "content": "run it"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": ""},
+                {
+                    "type": "tool_use",
+                    "id": "functions.Bash:0",
+                    "name": "Bash",
+                    "input": {},
+                    "provider_specific_fields": {"x": 1},
+                },
+            ],
+        },
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]},
+    ]
+
+    await call(server, messages=history)
+
+    body, _ = sent(server)
+    assert body["messages"] == [
+        {"role": "user", "content": "run it"},
+        {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]},
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(server: RecordingServer) -> None:
+    await call(
+        server,
+        provider_specific_header=[
+            {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}},
+            {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}},
+        ],
+    )
+
+    _, headers = sent(server)
+    assert headers["x-scoped"] == "yes"
+    assert "x-other" not in headers
+
+
+@pytest.mark.asyncio
+async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers(server: RecordingServer) -> None:
+    await call(
+        server,
+        headers={"x-priority": "forwarded", "x-forwarded-only": "kept"},
+        extra_headers={"x-priority": "extra", "x-extra-only": "kept"},
+        provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}},
+    )
+
+    _, headers = sent(server)
+    assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == {
+        "x-priority": "scoped",
+        "x-forwarded-only": "kept",
+        "x-extra-only": "kept",
+    }
+
+
+@pytest.mark.asyncio
+async def test_model_capabilities_are_read_from_model_cost_at_call_time(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    server.expected_requests = 2
+    model: Final = f"anthropic/{CAPABILITY_MODEL}"
+
+    capabilities(monkeypatch, sampling=True)
+    await call(server, model=model, temperature=0.2, drop_params=True)
+    capabilities(monkeypatch, sampling=False)
+    await call(server, model=model, temperature=0.2, drop_params=True)
+
+    assert_served_natively(server, count=2)
+    first, second = (request.body for request in server.requests)
+    assert isinstance(first, dict)
+    assert first["temperature"] == 0.2
+    assert second == {name: value for name, value in first.items() if name != "temperature"}
+
+
+@pytest.mark.asyncio
+async def test_global_drop_params_drops_unsupported_params(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    capabilities(monkeypatch, sampling=False)
+    monkeypatch.setattr(litellm, "drop_params", True)
+
+    await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5)
+
+    body, _ = sent(server)
+    assert "top_k" not in body
+
+
+@pytest.mark.asyncio
+async def test_unsupported_params_are_rejected_before_io_without_drop_params(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    capabilities(monkeypatch, sampling=False)
+    monkeypatch.setattr(litellm, "drop_params", False)
+    server.expected_requests = 0
+
+    with pytest.raises(litellm.BadRequestError, match="top_k"):
+        await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5)
+
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+async def test_optional_positional_body_parameters_reach_the_wire(server: RecordingServer) -> None:
+    metadata: Final = {"user_id": "caller"}
+    args: Final = (16, list(MESSAGES), "anthropic/claude-test", metadata, None, False, "Be brief", 0.25)
+
+    with native_route(Route.MESSAGES):
+        response: Final = await litellm.anthropic_messages(*args, api_key="test-key", api_base=server.base_url)
+
+    assert isinstance(response, dict)
+    assert response["id"] == MESSAGES_RESPONSE["id"]
+    body, _ = sent(server)
+    assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7])

tests/test_litellm_rust/messages/test_streaming.py

@@ -0,0 +1,73 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
+    server: RecordingServer,
+) -> None:
+    server.enqueue(STREAM)
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, stream=True, callbacks=[recorder])
+    assert isinstance(stream, AsyncIterator)
+    first: Final = await anext(stream)
+    await drain_logging()
+    assert "async_log_success_event" not in recorder.names
+    rest: Final = [chunk async for chunk in stream]
+
+    assert_served_natively(server)
+    assert first + b"".join(rest) == b"".join(STREAM.payloads())
+    success: Final = await recorder.wait_for_async("async_log_success_event")
+    assert len(success) == 1
+    assert isinstance(success[0].kwargs, dict)
+    assert success[0].kwargs["completion_start_time"] is not None
+    usage: Final = success[0].response.usage
+    assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
+    assert usage.completion_tokens == MESSAGES_RESPONSE["usage"]["output_tokens"]
+    assert "log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None:
+    server.enqueue(STREAM)
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder])
+    chunks: Final = await consume(stream)
+
+    assert_served_natively(server)
+    assert b"".join(chunks) == b"".join(STREAM.payloads())
+    assert len(await recorder.wait_for_async("async_log_success_event")) == 1
+
+
+@pytest.mark.asyncio
+async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> None:
+    server.enqueue(STREAM)
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, stream=True, callbacks=[recorder])
+    assert isinstance(stream, AsyncIterator)
+    await anext(stream)
+    await stream.aclose()
+
+    assert len(await recorder.wait_for_async("async_log_success_event")) == 1
+    with pytest.raises(StopAsyncIteration):
+        await anext(stream)

tests/test_litellm_rust/ocr/conftest.py

@@ -0,0 +1,19 @@
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.ocr.data import OCR_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
+    return recording_server
+
+
+@pytest.fixture(params=[False, True], ids=["python", "rust"])
+def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool:
+    enabled: Final = bool(request.param)
+    monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
+    return enabled

tests/test_litellm_rust/ocr/data.py

@@ -0,0 +1,39 @@
+from typing import Final
+
+import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
+OCR_MODEL: Final = "mistral/mistral-ocr-latest"
+OCR_TEXT: Final = "native OCR response"
+OCR_RESPONSE: Final = {
+    "pages": [{"index": 0, "markdown": OCR_TEXT, "images": [], "dimensions": None}],
+    "model": "mistral-ocr-latest",
+    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
+}
+
+
+def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
+    return {
+        "model": OCR_MODEL,
+        "document": dict(OCR_DOCUMENT),
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        **kwargs,
+    }
+
+
+def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
+    if not isinstance(response, OCRResponse):
+        raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
+    return response
+
+
+async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+    return await litellm.aocr(**ocr_arguments(server, **kwargs))
+
+
+async def call(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse:
+    return await call_aocr(server, **kwargs) if asynchronous else call_ocr(server, **kwargs)

tests/test_litellm_rust/ocr/test_callbacks.py

@@ -15,49 +15,38 @@ import litellm
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
-from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import (
+from tests.test_litellm_rust.ocr.data import (
     OCR_DOCUMENT,
     OCR_RESPONSE,
-    call_native,
-    call_native_aocr,
-    call_native_ocr,
+    call,
+    call_aocr,
+    call_ocr,
+)
+from tests.test_litellm_rust.support.callback_recorder import (
+    RecordingLogger,
+    drain_logging,
     request_body,
     request_headers,
 )
+from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
-    return recording_server
-
-
-def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
-    return call_native_ocr(server, callbacks=callbacks, **kwargs)
-
-
-async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
-    return await call_native_aocr(server, callbacks=callbacks, **kwargs)
-
-
-def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_callback_receives_transformed_provider_request(server: RecordingServer) -> None:
     observations: Final = []
 
     class Observe(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
             observations.append((model, copy.deepcopy(kwargs["additional_args"])))
 
-    call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0])
+    call_ocr(server, callbacks=[Observe()], pages=[0])
 
     assert len(observations) == 1
     model, additional_args = observations[0]
     assert model == "mistral-ocr-latest"
-    assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr"
+    assert additional_args["api_base"] == f"{server.base_url}/v1/ocr"
     assert additional_args["complete_input_dict"] == {
         "model": "mistral-ocr-latest",
         "document": OCR_DOCUMENT,
@@ -66,7 +55,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_
 
 
 @pytest.mark.asyncio
-async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(ocr_server: RecordingServer) -> None:
+async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(server: RecordingServer) -> None:
     class DropKeywords(CustomLogger):
         async def async_pre_call_deployment_hook(
             self, kwargs: dict[str, object], call_type: object
@@ -74,17 +63,17 @@ async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(o
             return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")}
 
     with rebound(litellm, "callbacks", [DropKeywords()]):
-        await call_native_aocr(ocr_server, pages=[0], extra_headers={"x-deleted": "remove this header"})
+        await call_aocr(server, pages=[0], extra_headers={"x-deleted": "remove this header"})
 
-    assert len(ocr_server.requests) == 1
-    sent: Final = ocr_server.requests[0]
+    assert len(server.requests) == 1
+    sent: Final = server.requests[0]
     assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
     assert sent.headers.get("x-deleted") is None
 
 
 @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
+    server: RecordingServer, raise_after_edit: bool
 ) -> None:
     observed: Final = []
 
@@ -98,13 +87,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
         def log_pre_api_call(self, model, messages, kwargs):
             observed.append(copy.deepcopy(request_body(kwargs)))
 
-    call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False)
+    call_ocr(server, callbacks=[Edit(), Observe()], include_image_base64=False)
 
     assert observed[0]["include_image_base64"] is True
-    assert ocr_server.requests[0].body["include_image_base64"] is True
+    assert server.requests[0].body["include_image_base64"] is True
 
 
-def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(server: RecordingServer) -> None:
     observed: Final = []
 
     class Edit(CustomLogger):
@@ -115,16 +104,16 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_
         def log_pre_api_call(self, model, messages, kwargs):
             observed.append(dict(request_headers(kwargs)))
 
-    call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()])
+    call_ocr(server, callbacks=[Edit(), Observe()])
 
     assert observed[0]["x-audit-tag"] == "reviewed"
-    assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed"
+    assert server.requests[0].headers["x-audit-tag"] == "reviewed"
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
-    ocr_server: RecordingServer, asynchronous: bool
+    server: RecordingServer, asynchronous: bool
 ) -> None:
     original: Final = dict(OCR_DOCUMENT)
     replacement_url: Final = "data:application/pdf;base64,ZGVm"
@@ -144,22 +133,20 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_
         "model": "mistral/mistral-ocr-latest",
         "document": original,
         "api_key": "test-key",
-        "api_base": ocr_server.base_url,
+        "api_base": server.base_url,
         "callbacks": [Retain(), Edit()],
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
 
     assert aliases == [True]
     assert retained[0]["document_url"] == replacement_url
     assert original["document_url"] == replacement_url
-    assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
+    assert server.requests[0].body["document"]["document_url"] == replacement_url
     assert response.pages[0].markdown == "native OCR response"
 
 
 def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     observed: Final = []
 
@@ -171,13 +158,13 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov
         def log_pre_api_call(self, model, messages, kwargs):
             observed.append(request_body(kwargs))
 
-    call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()])
+    call_ocr(server, callbacks=[Rebind(), Observe()])
 
     assert observed == [{"replacement": True}]
-    assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
+    assert server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
 
 
-def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None:
+def test_native_ocr_callback_retained_body_observes_later_callback_mutation(server: RecordingServer) -> None:
     queued: Final = []
 
     class QueuePayload(CustomLogger):
@@ -188,12 +175,12 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_
         def log_pre_api_call(self, model, messages, kwargs):
             request_body(kwargs)["queued-edit"] = True
 
-    call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()])
+    call_ocr(server, callbacks=[QueuePayload(), Edit()])
 
     assert queued[0]["queued-edit"] is True
 
 
-def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None:
+def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(server: RecordingServer) -> None:
     token: Final = object()
     terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue()
     finished: Final = threading.Event()
@@ -206,7 +193,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
             terminal_tokens.put(kwargs["test-token"])
             finished.set()
 
-    call_native_ocr_with_callbacks(ocr_server, [Stash()])
+    call_ocr(server, callbacks=[Stash()])
 
     assert finished.wait(10)
     assert terminal_tokens.get_nowait() is token
@@ -214,13 +201,13 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
 
 @pytest.mark.asyncio
 async def test_native_aocr_success_callback_receives_call_id_metadata_and_response(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     recorder: Final = RecordingLogger()
 
-    await call_native_aocr_with_callbacks(
-        ocr_server,
-        [recorder],
+    await call_aocr(
+        server,
+        callbacks=[recorder],
         litellm_call_id="ocr-success",
         metadata={"source": "callback-test"},
     )
@@ -235,9 +222,9 @@ async def test_native_aocr_success_callback_receives_call_id_metadata_and_respon
 
 @pytest.mark.asyncio
 async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
-    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
     observations: Final = []
 
     class Observe(CustomLogger):
@@ -248,7 +235,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
             observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
 
     with pytest.raises(litellm.InternalServerError):
-        await call_native_aocr_with_callbacks(ocr_server, [Observe()])
+        await call_aocr(server, callbacks=[Observe()])
 
     assert [observation[0] for observation in observations] == ["sync", "async"]
     assert all(observation[1] == "aocr" for observation in observations)
@@ -257,12 +244,12 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
 
 
 @pytest.mark.asyncio
-async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(server: RecordingServer) -> None:
     caller_loop: Final = asyncio.get_running_loop()
     caller_thread: Final = threading.current_thread()
     recorder: Final = RecordingLogger()
 
-    await call_native_aocr_with_callbacks(ocr_server, [recorder])
+    await call_aocr(server, callbacks=[recorder])
 
     events: Final = await recorder.wait_for_async("log_pre_api_call")
     assert len(events) == 1
@@ -272,9 +259,9 @@ async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_
 
 @pytest.mark.asyncio
 async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
-    ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+    server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
     token: Final = object()
     observed: Final = []
 
@@ -289,20 +276,20 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal
             observed.append(("async", kwargs["request-token"]))
 
     with pytest.raises(litellm.InternalServerError):
-        await call_native_aocr_with_callbacks(ocr_server, [TrackInFlightRequest()])
+        await call_aocr(server, callbacks=[TrackInFlightRequest()])
 
     assert [event for event, _ in observed] == ["sync", "async"]
     assert all(observed_token is token for _, observed_token in observed)
 
 
 def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     recorder: Final = RecordingLogger()
 
-    call_native_ocr_with_callbacks(
-        ocr_server,
-        [recorder, recorder],
+    call_ocr(
+        server,
+        callbacks=[recorder, recorder],
         success_callback=[recorder],
         failure_callback=[recorder],
     )
@@ -340,21 +327,21 @@ class ApplyEdits(CustomLogger):
 @settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
 @given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3))
 def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it(
-    ocr_server: RecordingServer, edits: dict[str, object]
+    server: RecordingServer, edits: dict[str, object]
 ) -> None:
-    ocr_server.expected_requests = None
+    server.expected_requests = None
 
     with isolated_callback_registries():
-        call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))])
+        call_ocr(server, callbacks=[ApplyEdits(MappingProxyType(edits))])
 
-    assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
+    assert server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
 
 
 @pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"])
-def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None:
+def test_native_ocr_sync_hooks_see_no_running_event_loop(server: RecordingServer, hook: str) -> None:
     recorder: Final = RecordingLogger()
 
-    call_native_ocr_with_callbacks(ocr_server, [recorder])
+    call_ocr(server, callbacks=[recorder])
 
     [event] = recorder.wait_for(hook)
     assert event.loop is None
@@ -364,7 +351,7 @@ def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingSe
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
-    ocr_server: RecordingServer, asynchronous: bool
+    server: RecordingServer, asynchronous: bool
 ) -> None:
     retained: Final = []
 
@@ -372,21 +359,21 @@ async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
         def log_pre_api_call(self, model, messages, kwargs):
             retained.append((kwargs, request_body(kwargs), request_headers(kwargs)))
 
-    await call_native(ocr_server, asynchronous, callbacks=[Retain()])
+    await call(server, asynchronous, callbacks=[Retain()])
     await drain_logging()
     gc.collect()
 
     [(details, body, headers)] = retained
-    assert body == ocr_server.requests[0].body
+    assert body == server.requests[0].body
     assert headers
-    assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items())
+    assert all(server.requests[0].headers[name.lower()] == value for name, value in headers.items())
     assert details["additional_args"]["complete_input_dict"] is body
     assert details["additional_args"]["headers"] is headers
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("family", ["sync", "async"])
-async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None:
+async def test_native_ocr_success_callbacks_share_one_logging_payload(server: RecordingServer, family: str) -> None:
     queued: Final = []
     finished: Final = threading.Event()
 
@@ -421,7 +408,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server
             async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
                 strip_payload(kwargs)
 
-    await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()])
+    await call(server, family == "async", callbacks=[QueuePayload(), StripPayload()])
     await drain_logging()
 
     assert await asyncio.to_thread(finished.wait, 10)
@@ -430,7 +417,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server
 
 @pytest.mark.asyncio
 async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     token: Final = object()
     observed: Final = []
@@ -455,7 +442,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
     litellm.callbacks.append(Block())
 
     with pytest.raises(Blocked) as raised:
-        await call_native_aocr(ocr_server)
+        await call_aocr(server)
     await drain_logging()
 
     assert observed == [("sync", token, raised.value), ("async", token, raised.value)]
@@ -464,7 +451,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
 ) -> None:
@@ -497,27 +484,25 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context
         "azure_ad_token_provider": provider,
         "callbacks": [Edit()],
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
     assert response.pages[0].markdown == "native OCR response"
     assert observations == ["token", "pre_call"]
-    assert ocr_server.requests[0].headers["authorization"] == "Bearer edited"
+    assert server.requests[0].headers["authorization"] == "Bearer edited"
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     isolated_azure_auth: None,
     asynchronous: bool,
 ) -> None:
-    ocr_server.expected_requests = 2
+    server.expected_requests = 2
     calls: Final = []
 
     def provider() -> str:
         calls.append("token")
-        nested: Final = call_native_ocr(ocr_server)
+        nested: Final = call_ocr(server)
         assert nested.pages[0].markdown == "native OCR response"
         return "outer-token"
 
@@ -526,12 +511,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
         "api_key": None,
         "azure_ad_token_provider": provider,
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
     assert response.pages[0].markdown == "native OCR response"
     assert calls == ["token"]
-    assert [request.headers["authorization"] for request in ocr_server.requests] == [
+    assert [request.headers["authorization"] for request in server.requests] == [
         "Bearer test-key",
         "Bearer outer-token",
     ]
@@ -539,10 +522,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
 
 @pytest.mark.asyncio
 async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     isolated_azure_auth: None,
 ) -> None:
-    ocr_server.expected_requests = 2
+    server.expected_requests = 2
 
     async def request(token: str, fail: bool) -> object:
         def provider() -> str:
@@ -550,8 +533,8 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
                 raise ValueError(token)
             return token
 
-        return await call_native_aocr(
-            ocr_server,
+        return await call_aocr(
+            server,
             model="azure_ai/mistral-ocr-latest",
             api_key=None,
             azure_ad_token_provider=provider,
@@ -567,7 +550,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
     assert isinstance(responses[1], litellm.APIConnectionError)
     assert "Failed to get Azure AD token: failed" in str(responses[1])
     assert isinstance(responses[2], OCRResponse)
-    assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [
+    assert sorted(request.headers["authorization"] for request in server.requests) == [
         "Bearer first",
         "Bearer second",
     ]
@@ -575,7 +558,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
 
 @pytest.mark.asyncio
 async def test_native_azure_ocr_releases_token_provider_after_cancellation(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     isolated_azure_auth: None,
 ) -> None:
     import gc
@@ -590,16 +573,16 @@ async def test_native_azure_ocr_releases_token_provider_after_cancellation(
     async def invoke() -> weakref.ReferenceType[Provider]:
         provider: Final = Provider()
         reference: Final = weakref.ref(provider)
-        ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
+        server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
         task: Final = asyncio.create_task(
-            call_native_aocr(
-                ocr_server,
+            call_aocr(
+                server,
                 model="azure_ai/mistral-ocr-latest",
                 api_key=None,
                 azure_ad_token_provider=provider,
             )
         )
-        await ocr_server.wait_for_requests(1)
+        await server.wait_for_requests(1)
         assert reference() is provider
         task.cancel()
         with pytest.raises(asyncio.CancelledError):

tests/test_litellm_rust/ocr/test_cohere.py

@@ -1,32 +0,0 @@
-from typing import Final
-
-import pytest
-
-import litellm
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-
-pytestmark = pytest.mark.requires_rust_extension
-MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")
-IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
-BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32}
-PAYLOAD: Final = {
-    "pages": [
-        {
-            "index": 4,
-            "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]},
-        },
-        {"markdown": {"content": "page two"}},
-    ],
-    "meta": {"billed_units": {"pages": 3}},
-}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model", MODELS)
-async def test_public_cohere_health_check(recording_server: RecordingServer, model: str) -> None:
-    recording_server.enqueue(ResponseSpec(body=PAYLOAD))
-    response: Final = await litellm.ahealth_check(
-        model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr"
-    )
-    assert "error" not in response
-    assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")

tests/test_litellm_rust/ocr/test_dispatch.py

@@ -1,42 +1,36 @@
-from typing import Final
+from typing import Final, NoReturn
 
 import pytest
 
 import litellm
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.routing import assert_served_natively
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
-    return recording_server
-
-
 @pytest.mark.parametrize("enabled", [False, True, None])
-def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: RecordingServer, enabled: bool | None) -> None:
+def test_public_ocr_uses_native_route_independently_of_flag(server: RecordingServer, enabled: bool | None) -> None:
     litellm.rust(enabled)
     response: Final = litellm.ocr(
         model=OCR_MODEL,
         document=OCR_DOCUMENT,
         api_key="test-key",
-        api_base=ocr_server.base_url,
+        api_base=server.base_url,
     )
 
     assert isinstance(response, OCRResponse)
     assert response.pages[0].markdown == "native OCR response"
-    assert_served_natively(ocr_server, response)
+    assert_served_natively(server, response)
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True])
 @pytest.mark.parametrize("caching", [None, False, True])
 async def test_ocr_does_not_depend_on_chat_cache(
-    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None
 ) -> None:
     from litellm.caching.caching import Cache
 
@@ -45,9 +39,93 @@ async def test_ocr_does_not_depend_on_chat_cache(
         "model": OCR_MODEL,
         "document": OCR_DOCUMENT,
         "api_key": "test-key",
-        "api_base": ocr_server.base_url,
+        "api_base": server.base_url,
         "caching": caching,
     }
     response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
     assert response.pages[0].markdown == "native OCR response"
-    assert len(ocr_server.requests) == 1
+    assert len(server.requests) == 1
+
+
+@pytest.mark.parametrize("required", [False, True])
+@pytest.mark.parametrize("failure", ["invalid", "live", "schema"])
+def test_native_projection_errors_never_select_python(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str
+) -> None:
+    import dataclasses
+    import ssl
+
+    from litellm.rust_bridge import runtime
+    from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
+    from litellm.rust_bridge.configuration import Rollout
+    from litellm.rust_bridge.host import settings
+    from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR
+    from litellm.rust_bridge.public_call import NativeCall
+
+    server.expected_requests = 0
+    snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1)
+    if failure == "schema":
+        monkeypatch.setattr(settings, "http_settings", lambda: snapshot)
+    else:
+        monkeypatch.setattr(
+            litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object()
+        )
+    request: Final = NativeCall(
+        args=(),
+        kwargs={},
+        base={
+            "model": "mistral/mistral-ocr-latest",
+            "document": OCR_DOCUMENT,
+            "api_key": "test-key",
+            "api_base": server.base_url,
+            "timeout": None,
+            "custom_llm_provider": "mistral",
+            "extra_headers": None,
+        },
+    )
+
+    def python_fallback() -> NoReturn:
+        pytest.fail("projection failures must not select Python")
+
+    with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"):
+        runtime.run(
+            RouteContext(Route.OCR, provider="mistral"),
+            binding=NATIVE_OCR,
+            native=lambda native: native(request),
+            python=python_fallback,
+            rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),),
+        )
+    assert server.requests == []
+
+
+MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")
+
+
+IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
+
+
+BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32}
+
+
+PAYLOAD: Final = {
+    "pages": [
+        {
+            "index": 4,
+            "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]},
+        },
+        {"markdown": {"content": "page two"}},
+    ],
+    "meta": {"billed_units": {"pages": 3}},
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model", MODELS)
+async def test_ocr_health_check_uses_the_native_route(recording_server: RecordingServer, model: str) -> None:
+    recording_server.enqueue(ResponseSpec(body=PAYLOAD))
+    response: Final = await litellm.ahealth_check(
+        model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr"
+    )
+    assert "error" not in response
+    assert_served_natively(recording_server)
+    assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")

tests/test_litellm_rust/ocr/test_failures.py

@@ -0,0 +1,340 @@
+import json
+from collections.abc import Callable
+from dataclasses import dataclass
+from io import BytesIO
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+from pydantic import JsonValue
+
+import litellm
+from tests.test_litellm_rust.ocr.data import (
+    OCR_RESPONSE,
+    call,
+    call_aocr,
+    call_ocr,
+)
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_ocr_contract_upstream_status(
+    server: RecordingServer,
+    ocr_backend: bool,
+    asynchronous: bool,
+) -> None:
+    upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422)
+    server.enqueue(upstream)
+    arguments: Final = {
+        "model": "vertex_ai/mistral-ocr-latest",
+        "vertex_project": "test-project",
+        "vertex_location": "us-central1",
+        "num_retries": 0,
+    }
+    with pytest.raises(litellm.BadRequestError) as caught:
+        await call(server, asynchronous, **arguments)
+    assert caught.value.status_code == upstream.status
+    assert caught.value.response.status_code == upstream.status
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("preserved", ["body", "headers"])
+async def test_ocr_contract_provider_error_details(
+    server: RecordingServer,
+    ocr_backend: bool,
+    asynchronous: bool,
+    preserved: str,
+) -> None:
+    payload: Final = {"message": "rate limited"}
+    headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"}
+    server.enqueue(ResponseSpec(body=payload, status=429, headers=headers))
+    with pytest.raises(litellm.RateLimitError) as caught:
+        await call(server, asynchronous, num_retries=0)
+    response: Final = caught.value.response
+    assert isinstance(response, httpx.Response)
+    if preserved == "body":
+        assert response.content == json.dumps(payload).encode()
+    else:
+        for name, value in headers.items():
+            assert response.headers.get(name.lower()) == value
+            assert response.headers.get(name.upper()) == value
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_ocr_contract_invalid_response_format(
+    server: RecordingServer,
+    ocr_backend: bool,
+    asynchronous: bool,
+) -> None:
+    server.expected_requests = 0
+    with pytest.raises(litellm.UnsupportedParamsError) as caught:
+        await call(server, asynchronous, req_format="bogus", num_retries=0)
+    assert caught.value.status_code == 400
+    for value in ("req_format", "bogus", "native", "litellm"):
+        assert value in str(caught.value)
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    "document,field",
+    [
+        ([], "document"),
+        ({"document_url": "https://example.com/a.pdf"}, "type"),
+        ({"type": "text"}, "type"),
+    ],
+)
+async def test_ocr_contract_malformed_document_is_actionable(
+    server: RecordingServer,
+    ocr_backend: bool,
+    asynchronous: bool,
+    document: JsonValue,
+    field: str,
+) -> None:
+    server.expected_requests = None
+    with pytest.raises(litellm.BadRequestError) as caught:
+        await call(server, asynchronous, document=document, num_retries=0)
+    assert caught.value.status_code == 400
+    assert field.lower() in str(caught.value).lower()
+    assert "NoneType: None" not in str(caught.value)
+    assert "indices must be" not in str(caught.value)
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")])
+async def test_ocr_contract_azure_invalid_options_are_bad_requests(
+    server: RecordingServer,
+    ocr_backend: bool,
+    asynchronous: bool,
+    option: str,
+    value: JsonValue,
+    field: str,
+) -> None:
+    server.expected_requests = 0
+    arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0}
+    with pytest.raises(litellm.BadRequestError) as caught:
+        await call(server, asynchronous, **arguments)
+    assert caught.value.status_code == 400
+    assert field in str(caught.value)
+    assert server.requests == []
+
+
+def test_native_ocr_maps_provider_400_with_public_provider_details(server: RecordingServer) -> None:
+    server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
+
+    with pytest.raises(litellm.BadRequestError) as caught:
+        call_ocr(server)
+
+    assert caught.value.status_code == 400
+    assert caught.value.model == "mistral-ocr-latest"
+    assert caught.value.llm_provider == "mistral"
+    assert "invalid OCR request" in str(caught.value)
+
+
+class TokenAbort(BaseException):
+    pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"])
+async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
+    server: RecordingServer,
+    isolated_azure_auth: None,
+    asynchronous: bool,
+    failure: str,
+) -> None:
+    server.expected_requests = 0
+    calls: Final = []
+    recorder: Final = RecordingLogger()
+    original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")}
+
+    def token_provider() -> object:
+        calls.append("token")
+        raise original[failure]
+
+    arguments: Final = {
+        "model": "azure_ai/mistral-ocr-latest",
+        "api_key": None,
+        "azure_ad_token_provider": token_provider,
+        "callbacks": [recorder],
+    }
+    expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
+    with pytest.raises(expected) as caught:
+        await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
+    assert calls == ["token"]
+    assert server.requests == []
+    assert "log_pre_api_call" not in recorder.names
+    if failure == "ordinary":
+        assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
+        assert isinstance(caught.value.__context__, RuntimeError)
+        assert caught.value.__context__.__cause__ is original[failure]
+    else:
+        assert caught.value is original[failure]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_native_file_preparation_preserves_reader_exception(server: RecordingServer, asynchronous: bool) -> None:
+    server.expected_requests = 0
+    failure: Final = RuntimeError("reader failed")
+
+    class Reader:
+        def read(self) -> bytes:
+            raise failure
+
+    document: Final = {"type": "file", "file": Reader()}
+    with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
+        await call_aocr(server, document=document) if asynchronous else call_ocr(server, document=document)
+    assert caught.value.__context__ is failure
+
+
+COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
+
+
+FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024
+
+
+class IntReader:
+    def read(self) -> int:
+        return 1
+
+
+def oversized_file(tmp_path: Path) -> Path:
+    path: Final = tmp_path / "large.pdf"
+    with path.open("wb") as stream:
+        stream.truncate(FILE_SIZE_LIMIT + 1)
+    return path
+
+
+def empty_token() -> str:
+    return ""
+
+
+def unused_token() -> str:
+    raise AssertionError("the token provider must not run")
+
+
+@dataclass(frozen=True, slots=True)
+class PublicFailure:
+    arguments: Callable[[Path], dict[str, object]]
+    error: type[Exception]
+    match: str
+    provider_requests: int = 0
+    response: ResponseSpec | None = None
+    cause: type[BaseException] | None = None
+
+
+PUBLIC_FAILURES: Final = {
+    "unknown-req-format": PublicFailure(
+        lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`"
+    ),
+    "empty-file": PublicFailure(
+        lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty"
+    ),
+    "oversized-file": PublicFailure(
+        lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}},
+        litellm.BadRequestError,
+        "exceeds the size limit",
+    ),
+    "missing-file": PublicFailure(
+        lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}},
+        litellm.APIConnectionError,
+        "File not found",
+        cause=FileNotFoundError,
+    ),
+    "reader-returns-non-bytes": PublicFailure(
+        lambda _: {"document": {"type": "file", "file": IntReader()}},
+        litellm.APIConnectionError,
+        "bytes or str",
+        cause=TypeError,
+    ),
+    "cohere-non-image": PublicFailure(
+        lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`"
+    ),
+    "cohere-unknown-format": PublicFailure(
+        lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"},
+        litellm.BadRequestError,
+        "output_format",
+    ),
+    "azure-missing-api-base": PublicFailure(
+        lambda _: {
+            "model": "azure_ai/mistral-ocr-latest",
+            "api_key": None,
+            "api_base": None,
+            "azure_ad_token_provider": unused_token,
+        },
+        litellm.APIConnectionError,
+        "Missing Azure AI API Base",
+    ),
+    "azure-empty-token": PublicFailure(
+        lambda _: {
+            "model": "azure_ai/mistral-ocr-latest",
+            "api_key": None,
+            "azure_ad_token": "static-token",
+            "azure_ad_token_provider": empty_token,
+        },
+        litellm.APIConnectionError,
+        "Missing Azure AI credentials",
+    ),
+    "upstream-500": PublicFailure(
+        lambda _: {},
+        litellm.InternalServerError,
+        "provider unavailable",
+        provider_requests=1,
+        response=ResponseSpec(body={"message": "provider unavailable"}, status=500),
+    ),
+    "invalid-provider-response": PublicFailure(
+        lambda _: {},
+        litellm.APIConnectionError,
+        "pages",
+        provider_requests=1,
+        response=ResponseSpec(body={"pages": "invalid"}),
+    ),
+    "response-over-limit": PublicFailure(
+        lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1},
+        litellm.APIConnectionError,
+        "OCR response exceeds the size limit",
+        provider_requests=1,
+    ),
+    "timeout": PublicFailure(
+        lambda _: {"timeout": 0.01},
+        litellm.Timeout,
+        "",
+        provider_requests=1,
+        response=ResponseSpec(body=OCR_RESPONSE, delay=0.2),
+    ),
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys())
+async def test_native_failures_raise_the_public_exception_class(
+    server: RecordingServer,
+    isolated_azure_auth: None,
+    tmp_path: Path,
+    asynchronous: bool,
+    failure: PublicFailure,
+) -> None:
+    server.expected_requests = failure.provider_requests
+    if failure.response is not None:
+        server.enqueue(failure.response)
+
+    with pytest.raises(failure.error, match=failure.match) as caught:
+        await call(server, asynchronous, **failure.arguments(tmp_path))
+
+    assert len(server.requests) == failure.provider_requests
+    if failure.cause is not None:
+        assert isinstance(caught.value.__context__, failure.cause)

tests/test_litellm_rust/ocr/test_guardrails.py

@@ -8,19 +8,13 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
 from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
 from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.ocr.data import call_aocr
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native, call_native_aocr
+from tests.test_litellm_rust.support.recording_server import RecordingServer
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
-    return recording_server
-
-
 class ReplaceOCRMarkdown(CustomGuardrail):
     def __init__(self) -> None:
         super().__init__(
@@ -36,7 +30,7 @@ class ReplaceOCRMarkdown(CustomGuardrail):
 
 @pytest.mark.asyncio
 async def test_native_aocr_post_call_content_filter_blocks_matching_markdown(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     guardrail: Final = ContentFilterGuardrail(
         guardrail_name="block-native-ocr-markdown",
@@ -46,22 +40,22 @@ async def test_native_aocr_post_call_content_filter_blocks_matching_markdown(
     litellm.callbacks.append(guardrail)
 
     with pytest.raises(HTTPException, match="Content blocked") as blocked:
-        await call_native_aocr(ocr_server, guardrails=[guardrail.guardrail_name])
+        await call_aocr(server, guardrails=[guardrail.guardrail_name])
 
     assert blocked.value.status_code == 400
-    assert len(ocr_server.requests) == 1
+    assert len(server.requests) == 1
 
 
 @pytest.mark.asyncio
 async def test_native_aocr_post_call_replacement_reaches_caller_and_success_callback(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     guardrail: Final = ReplaceOCRMarkdown()
     recorder: Final = RecordingLogger()
     litellm.callbacks.append(guardrail)
 
-    response: Final = await call_native_aocr(
-        ocr_server,
+    response: Final = await call_aocr(
+        server,
         callbacks=[recorder],
         guardrails=[guardrail.guardrail_name],
     )
@@ -71,4 +65,4 @@ async def test_native_aocr_post_call_replacement_reaches_caller_and_success_call
     assert response.pages[0].markdown == "Reviewed OCR"
     assert len(success_events) == 1
     assert success_events[0].response.pages[0].markdown == "Reviewed OCR"
-    assert "guardrails" not in ocr_server.requests[0].body
+    assert "guardrails" not in server.requests[0].body

tests/test_litellm_rust/ocr/test_lifecycle.py

@@ -15,51 +15,45 @@ from litellm._logging import trace_id_var
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from tests.test_litellm_rust.ocr.data import OCR_RESPONSE, call_aocr, call_ocr
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr, call_ocr
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
-    return recording_server
-
-
 @pytest.mark.asyncio
-async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) -> None:
+async def test_proxy_metadata_remains_python_owned(server: RecordingServer) -> None:
     from litellm.proxy._types import UserAPIKeyAuth
 
     recorder: Final = RecordingLogger()
     auth: Final = UserAPIKeyAuth(user_id="ocr-user")
     response: Final = await call_aocr(
-        ocr_server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object()
+        server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object()
     )
     events: Final = await recorder.wait_for_async("async_log_success_event")
     assert response.pages[0].markdown == "native OCR response"
     assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
     assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user"
-    assert "metadata" not in ocr_server.requests[0].body
+    assert "metadata" not in server.requests[0].body
 
 
 @pytest.mark.asyncio
 async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
 ) -> None:
     recorder: Final = RecordingLogger()
-    response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05)
+    response: Final = await call_aocr(server, callbacks=[recorder], ocr_cost_per_page=0.05)
     events: Final = await recorder.wait_for_async("async_log_success_event")
 
     assert response.usage_info is not None and response.usage_info.pages_processed == 1
     assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05
     assert response._hidden_params["response_cost"] == pytest.approx(0.05)
-    assert "ocr_cost_per_page" not in ocr_server.requests[0].body
+    assert "ocr_cost_per_page" not in server.requests[0].body
 
 
 @pytest.mark.asyncio
-async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None:
+async def test_response_replacement_finalized_before_dispatch_in_caller_task(server: RecordingServer) -> None:
     caller: Final = asyncio.current_task()
     context: Final = ContextVar("lifecycle-test", default="before")
     observations: Final = []
@@ -76,11 +70,11 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr
             return response.model_copy(update={"model": "replaced"})
 
     litellm.callbacks.append(Replace())
-    response: Final = await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="native-final")
+    response: Final = await call_aocr(server, callbacks=[recorder], litellm_call_id="native-final")
     events: Final = await recorder.wait_for_async("async_log_success_event")
     assert observations == [("pre", caller, "pre"), ("post", caller, "pre")]
     assert context.get() == "pre"
-    assert ocr_server.requests[0].body["pages"] == [2]
+    assert server.requests[0].body["pages"] == [2]
     assert response.model == "replaced"
     assert events[0].response is response
     assert response._hidden_params["litellm_call_id"] == "native-final"
@@ -90,7 +84,7 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True])
 async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
-    ocr_server: RecordingServer, asynchronous: bool
+    server: RecordingServer, asynchronous: bool
 ) -> None:
     failure: Final = RuntimeError("metadata failed")
     seen: Final = []
@@ -120,8 +114,8 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
         )
         reference: Final = weakref.ref(logger)
         with pytest.raises(RuntimeError) as caught:
-            await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr(
-                ocr_server, litellm_logging_obj=logger
+            await call_aocr(server, litellm_logging_obj=logger) if asynchronous else call_ocr(
+                server, litellm_logging_obj=logger
             )
         assert caught.value is failure
         failure.__traceback__ = None
@@ -132,12 +126,12 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
     gc.collect()
     assert seen == ([("sync", failure), ("async", failure)] if asynchronous else [("sync", failure)])
     assert reference() is None
-    assert len(ocr_server.requests) == 1
+    assert len(server.requests) == 1
 
 
 @pytest.mark.asyncio
-async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: RecordingServer) -> None:
-    ocr_server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500))
+async def test_mapped_failure_identity_and_deployment_snapshot(server: RecordingServer) -> None:
+    server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500))
     recorder: Final = RecordingLogger()
     snapshots: Final = []
 
@@ -148,20 +142,20 @@ async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: Recor
 
     litellm.callbacks.append(Observe())
     with pytest.raises(litellm.InternalServerError) as caught:
-        await call_aocr(ocr_server, callbacks=[recorder])
+        await call_aocr(server, callbacks=[recorder])
     failures: Final = tuple(event for event in recorder.events if "failure" in event.name)
     assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"]
     assert all(event.kwargs["exception"] is caught.value for event in failures)
     assert caught.value.status_code == 500
     assert snapshots[0] is not caught.value
     assert snapshots[0].status_code == 418
-    assert len(ocr_server.requests) == 1
+    assert len(server.requests) == 1
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("phase", ["pre", "http"])
 async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
-    ocr_server: RecordingServer, phase: str
+    server: RecordingServer, phase: str
 ) -> None:
     entered: Final = asyncio.Event()
     recorder: Final = RecordingLogger()
@@ -174,21 +168,21 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
 
     litellm.callbacks.append(Pause())
     if phase == "http":
-        ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
+        server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
     if phase == "pre":
-        ocr_server.expected_requests = 0
+        server.expected_requests = 0
     restored: Final = []
 
     async def invoke():
         trace_id_var.set("parent")
         try:
-            await call_aocr(ocr_server, callbacks=[recorder], litellm_trace_id="native-call")
+            await call_aocr(server, callbacks=[recorder], litellm_trace_id="native-call")
         finally:
             restored.append(trace_id_var.get())
 
     task: Final = asyncio.create_task(invoke())
     if phase == "http":
-        await ocr_server.wait_for_requests(1)
+        await server.wait_for_requests(1)
     else:
         await asyncio.wait_for(entered.wait(), 5)
     task.cancel()
@@ -201,9 +195,7 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("blocked", [False, True])
-async def test_deferred_logging_requires_release_and_runs_at_most_once(
-    ocr_server: RecordingServer, blocked: bool
-) -> None:
+async def test_deferred_logging_requires_release_and_runs_at_most_once(server: RecordingServer, blocked: bool) -> None:
     recorder: Final = RecordingLogger()
     logger: Final = Logging(
         model="mistral-ocr-latest",
@@ -216,7 +208,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(
         dynamic_async_success_callbacks=[recorder],
     )
     logger._defer_async_logging = True
-    response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger)
+    response: Final = await call_aocr(server, litellm_logging_obj=logger)
     await drain_logging()
     assert "async_log_success_event" not in recorder.names
     ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
@@ -228,7 +220,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(
         assert events[0].response is response
 
 
-def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: RecordingServer) -> None:
+def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None:
     context: Final = ContextVar("sync-lifecycle", default="missing")
     context.set("caller")
     thread: Final = threading.current_thread()
@@ -240,7 +232,7 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record
             observations.append((threading.current_thread(), context.get(), response_obj))
             finished.set()
 
-    response: Final = call_ocr(ocr_server, callbacks=[Observe()])
+    response: Final = call_ocr(server, callbacks=[Observe()])
     assert finished.wait(5)
     assert observations[0][0] is not thread
     assert observations[0][1] == "caller"
@@ -248,44 +240,44 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record
 
 
 @pytest.mark.asyncio
-async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(ocr_server: RecordingServer) -> None:
-    ocr_server.expected_requests = 2
+async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(server: RecordingServer) -> None:
+    server.expected_requests = 2
     recorder: Final = RecordingLogger()
     outcomes: Final = []
 
     class Nested(CustomLogger):
         async def async_pre_call_deployment_hook(self, kwargs, call_type):
             if kwargs.get("litellm_call_id") == "outer":
-                outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="inner"))
+                outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="inner"))
 
     litellm.callbacks.append(Nested())
-    outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="outer"))
+    outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="outer"))
     events: Final = await recorder.wait_for_async("async_log_success_event", count=2)
     assert [event.kwargs["litellm_call_id"] for event in events] == ["inner", "outer"]
     assert events[0].response is outcomes[0]
     assert events[1].response is outcomes[1]
-    assert len(ocr_server.requests) == 2
+    assert len(server.requests) == 2
 
 
-def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServer) -> None:
-    ocr_server.expected_requests = 2
+def test_sync_pre_call_can_make_nested_native_request(server: RecordingServer) -> None:
+    server.expected_requests = 2
     observed: Final = []
 
     class Nested(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
             if kwargs["litellm_call_id"] == "outer-sync":
-                observed.append(call_ocr(ocr_server, litellm_call_id="inner-sync"))
+                observed.append(call_ocr(server, litellm_call_id="inner-sync"))
 
-    response: Final = call_ocr(ocr_server, callbacks=[Nested()], litellm_call_id="outer-sync")
+    response: Final = call_ocr(server, callbacks=[Nested()], litellm_call_id="outer-sync")
     assert observed[0].pages[0].markdown == response.pages[0].markdown
-    assert len(ocr_server.requests) == 2
+    assert len(server.requests) == 2
 
 
-def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None:
+def test_unstarted_native_coroutine_releases_input_without_reading_file(server: RecordingServer) -> None:
     from litellm.ocr.dispatch import _public_request
     from litellm.rust_bridge import _native
 
-    ocr_server.expected_requests = 0
+    server.expected_requests = 0
     effects: Final = []
 
     class File:
@@ -308,7 +300,7 @@ def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_serv
 
 
 @pytest.mark.asyncio
-async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server: RecordingServer) -> None:
+async def test_file_read_happens_after_deployment_hook_in_caller_task(server: RecordingServer) -> None:
     effects: Final = []
     caller: Final = asyncio.current_task()
 
@@ -323,13 +315,13 @@ async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server
             effects.append(("hook", asyncio.current_task()))
 
     litellm.callbacks.append(Deployment())
-    await call_aocr(ocr_server, document={"type": "file", "file": File()})
+    await call_aocr(server, document={"type": "file", "file": File()})
     assert effects == [("hook", caller), ("read", caller)]
 
 
 @pytest.mark.asyncio
-async def test_failure_callbacks_continue_within_both_families(ocr_server: RecordingServer) -> None:
-    ocr_server.enqueue(ResponseSpec(body={"message": "failed"}, status=500))
+async def test_failure_callbacks_continue_within_both_families(server: RecordingServer) -> None:
+    server.enqueue(ResponseSpec(body={"message": "failed"}, status=500))
     observed: Final = []
 
     class Broken(CustomLogger):
@@ -349,7 +341,7 @@ async def test_failure_callbacks_continue_within_both_families(ocr_server: Recor
             observed.append(("following-async", kwargs["exception"]))
 
     with pytest.raises(litellm.InternalServerError) as caught:
-        await call_aocr(ocr_server, callbacks=[Broken(), Following()])
+        await call_aocr(server, callbacks=[Broken(), Following()])
     assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"]
     assert all(error is caught.value for _, error in observed)
 
@@ -392,17 +384,17 @@ async def test_cancelling_native_transport_closes_connection_before_return() ->
 
 
 @pytest.mark.asyncio
-async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: RecordingServer) -> None:
+async def test_response_limit_is_enforced_at_the_public_boundary(server: RecordingServer) -> None:
     limit: Final = len(json.dumps(OCR_RESPONSE).encode()) - 1
     with pytest.raises(litellm.APIConnectionError, match="OCR response exceeds the size limit"):
-        await call_aocr(ocr_server, max_response_bytes=limit)
-    assert len(ocr_server.requests) == 1
+        await call_aocr(server, max_response_bytes=limit)
+    assert len(server.requests) == 1
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("failure", [False, True])
 async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     monkeypatch: pytest.MonkeyPatch,
     failure: bool,
     created_loggers: list[Logging],
@@ -439,14 +431,14 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
     monkeypatch.setattr(litellm_logging, "executor", probe)
     monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe)
     if failure:
-        ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500))
+        server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500))
     trace_id_var.set("callback-free-parent")
     arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"}
     if failure:
         with pytest.raises(litellm.InternalServerError):
-            await call_aocr(ocr_server, **arguments)
+            await call_aocr(server, **arguments)
     else:
-        response: Final = await call_aocr(ocr_server, **arguments)
+        response: Final = await call_aocr(server, **arguments)
         assert response.pages[0].markdown == "native OCR response"
         assert response._hidden_params["litellm_call_id"] == "callback-free-id"
         assert response._hidden_params["response_cost"] is not None
@@ -468,20 +460,18 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
 @pytest.mark.parametrize(
     "registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"]
 )
-async def test_terminal_registration_added_during_http_is_observed(
-    ocr_server: RecordingServer, registration: str
-) -> None:
+async def test_terminal_registration_added_during_http_is_observed(server: RecordingServer, registration: str) -> None:
     failure: Final = "failure" in registration
     observer: Final = RecordingLogger()
-    ocr_server.enqueue(
+    server.enqueue(
         ResponseSpec(
             body={"message": "provider failed"} if failure else OCR_RESPONSE, status=500 if failure else 200, delay=0.1
         )
     )
     task: Final = asyncio.create_task(
-        asyncio.to_thread(call_ocr, ocr_server) if registration == "success_callback" else call_aocr(ocr_server)
+        asyncio.to_thread(call_ocr, server) if registration == "success_callback" else call_aocr(server)
     )
-    await ocr_server.wait_for_requests(1)
+    await server.wait_for_requests(1)
     getattr(litellm, registration).append(observer)
     if failure:
         with pytest.raises(litellm.InternalServerError):
@@ -523,7 +513,7 @@ def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]:
 @pytest.mark.asyncio
 @pytest.mark.parametrize("consumer", ["logger_fn", "raw_global", "request_debug"])
 async def test_explicit_logging_consumers_keep_request_and_response_payloads(
-    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str
 ) -> None:
     snapshots: Final[list[dict[str, object]]] = []
     if consumer == "raw_global":
@@ -533,7 +523,7 @@ async def test_explicit_logging_consumers_keep_request_and_response_payloads(
         "raw_global": {},
         "request_debug": {"litellm_request_debug": True},
     }[consumer]
-    response: Final = await call_aocr(ocr_server, **arguments)
+    response: Final = await call_aocr(server, **arguments)
     details: Final = created_loggers[0].model_call_details
     assert details["additional_args"]["complete_input_dict"]["model"] == "mistral-ocr-latest"
     assert json.loads(details["original_response"])["pages"][0]["markdown"] == response.pages[0].markdown

tests/test_litellm_rust/ocr/test_requests.py

@@ -1,160 +1,36 @@
-import json
-from collections.abc import Callable
-from dataclasses import dataclass
 from io import BytesIO
-from pathlib import Path
-from typing import Final, NoReturn
+from typing import Final
 
-import httpx
 import pytest
-from pydantic import JsonValue
 
-import litellm
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import (
+from tests.test_litellm_rust.ocr.data import (
     OCR_DOCUMENT,
     OCR_RESPONSE,
-    call_native,
-    call_native_aocr,
-    call_native_ocr,
+    call_aocr,
+    call_ocr,
 )
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+from tests.test_litellm_rust.support.routing import assert_served_natively
 
 pytestmark = pytest.mark.requires_rust_extension
 
 
-@pytest.fixture(params=[False, True], ids=["python", "rust"])
-def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool:
-    enabled: Final = bool(request.param)
-    monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
-    return enabled
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_ocr_contract_upstream_status(
-    ocr_server: RecordingServer,
-    ocr_backend: bool,
-    asynchronous: bool,
-) -> None:
-    upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422)
-    ocr_server.enqueue(upstream)
-    arguments: Final = {
-        "model": "vertex_ai/mistral-ocr-latest",
-        "vertex_project": "test-project",
-        "vertex_location": "us-central1",
-        "num_retries": 0,
-    }
-    with pytest.raises(litellm.BadRequestError) as caught:
-        await call_native(ocr_server, asynchronous, **arguments)
-    assert caught.value.status_code == upstream.status
-    assert caught.value.response.status_code == upstream.status
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("preserved", ["body", "headers"])
-async def test_ocr_contract_provider_error_details(
-    ocr_server: RecordingServer,
-    ocr_backend: bool,
-    asynchronous: bool,
-    preserved: str,
-) -> None:
-    payload: Final = {"message": "rate limited"}
-    headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"}
-    ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers))
-    with pytest.raises(litellm.RateLimitError) as caught:
-        await call_native(ocr_server, asynchronous, num_retries=0)
-    response: Final = caught.value.response
-    assert isinstance(response, httpx.Response)
-    if preserved == "body":
-        assert response.content == json.dumps(payload).encode()
-    else:
-        for name, value in headers.items():
-            assert response.headers.get(name.lower()) == value
-            assert response.headers.get(name.upper()) == value
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_ocr_contract_invalid_response_format(
-    ocr_server: RecordingServer,
-    ocr_backend: bool,
-    asynchronous: bool,
-) -> None:
-    ocr_server.expected_requests = 0
-    with pytest.raises(litellm.UnsupportedParamsError) as caught:
-        await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0)
-    assert caught.value.status_code == 400
-    for value in ("req_format", "bogus", "native", "litellm"):
-        assert value in str(caught.value)
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
-    "document,field",
-    [
-        ([], "document"),
-        ({"document_url": "https://example.com/a.pdf"}, "type"),
-        ({"type": "text"}, "type"),
-    ],
-)
-async def test_ocr_contract_malformed_document_is_actionable(
-    ocr_server: RecordingServer,
-    ocr_backend: bool,
-    asynchronous: bool,
-    document: JsonValue,
-    field: str,
-) -> None:
-    ocr_server.expected_requests = None
-    with pytest.raises(litellm.BadRequestError) as caught:
-        await call_native(ocr_server, asynchronous, document=document, num_retries=0)
-    assert caught.value.status_code == 400
-    assert field.lower() in str(caught.value).lower()
-    assert "NoneType: None" not in str(caught.value)
-    assert "indices must be" not in str(caught.value)
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")])
-async def test_ocr_contract_azure_invalid_options_are_bad_requests(
-    ocr_server: RecordingServer,
-    ocr_backend: bool,
-    asynchronous: bool,
-    option: str,
-    value: JsonValue,
-    field: str,
-) -> None:
-    ocr_server.expected_requests = 0
-    arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0}
-    with pytest.raises(litellm.BadRequestError) as caught:
-        await call_native(ocr_server, asynchronous, **arguments)
-    assert caught.value.status_code == 400
-    assert field in str(caught.value)
-    assert ocr_server.requests == []
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 @pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"])
 async def test_ocr_contract_native_format_supported(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     ocr_backend: bool,
     asynchronous: bool,
     model: str,
 ) -> None:
-    ocr_server.expected_requests = None
+    server.expected_requests = None
     payload: Final = (
         {"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}}
         if model.startswith("reducto/")
         else OCR_RESPONSE
     )
-    ocr_server.default_response = ResponseSpec(body=payload)
+    server.default_response = ResponseSpec(body=payload)
     arguments: Final = {
         "model": model,
         "req_format": "native",
@@ -163,43 +39,39 @@ async def test_ocr_contract_native_format_supported(
         if model.startswith("reducto/")
         else OCR_DOCUMENT,
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
     assert response.pages[0].markdown == "native OCR response"
     assert response.get_provider_native_response() == payload
-    assert len(ocr_server.requests) == 1
+    assert len(server.requests) == 1
     if ocr_backend:
-        assert_served_natively(ocr_server, response)
+        assert_served_natively(server, response)
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_ocr_contract_unknown_reducto_model_reaches_provider(
-    ocr_server: RecordingServer,
+    server: RecordingServer,
     ocr_backend: bool,
     asynchronous: bool,
 ) -> None:
-    ocr_server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}})
+    server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}})
     arguments: Final = {
         "model": "reducto/future-parse-model",
         "document": {"type": "document_url", "document_url": "reducto://ready.pdf"},
         "num_retries": 0,
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
     assert response.model == "future-parse-model"
     assert response.pages[0].markdown == "future model response"
-    assert len(ocr_server.requests) == 1
-    assert ocr_server.requests[0].path == "/parse"
-    assert ocr_server.requests[0].body == {"input": "reducto://ready.pdf"}
+    assert len(server.requests) == 1
+    assert server.requests[0].path == "/parse"
+    assert server.requests[0].body == {"input": "reducto://ready.pdf"}
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
 async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token(
-    ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
+    server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
 ) -> None:
     calls: Final = []
 
@@ -212,478 +84,24 @@ async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token(
         "api_key": None,
         "azure_ad_token_provider": token_provider,
     }
-    response: Final = (
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    )
+    response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
 
     assert calls == ["token"]
     assert response.pages[0].markdown == "native OCR response"
-    assert_served_natively(ocr_server, response)
-    assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token"
+    assert_served_natively(server, response)
+    assert server.requests[0].headers["authorization"] == "Bearer callback-token"
 
 
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
-    recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
-    return recording_server
-
-
-def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None:
-    ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
-
-    with pytest.raises(litellm.BadRequestError) as caught:
-        call_native_ocr(ocr_server)
-
-    assert caught.value.status_code == 400
-    assert caught.value.model == "mistral-ocr-latest"
-    assert caught.value.llm_provider == "mistral"
-    assert "invalid OCR request" in str(caught.value)
-
-
-def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_server: RecordingServer) -> None:
-    response: Final = call_native_ocr(
-        ocr_server,
+def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(server: RecordingServer) -> None:
+    response: Final = call_ocr(
+        server,
         document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"},
         opaque_extension=object(),
     )
 
     assert response.pages[0].markdown == "native OCR response"
-    assert_served_natively(ocr_server, response)
-    assert ocr_server.requests[0].body == {
+    assert_served_natively(server, response)
+    assert server.requests[0].body == {
         "model": "mistral-ocr-latest",
         "document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"},
     }
-
-
-class TokenAbort(BaseException):
-    pass
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"])
-async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
-    ocr_server: RecordingServer,
-    isolated_azure_auth: None,
-    asynchronous: bool,
-    failure: str,
-) -> None:
-    ocr_server.expected_requests = 0
-    calls: Final = []
-    recorder: Final = RecordingLogger()
-    original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")}
-
-    def token_provider() -> object:
-        calls.append("token")
-        raise original[failure]
-
-    arguments: Final = {
-        "model": "azure_ai/mistral-ocr-latest",
-        "api_key": None,
-        "azure_ad_token_provider": token_provider,
-        "callbacks": [recorder],
-    }
-    expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
-    with pytest.raises(expected) as caught:
-        await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
-    assert calls == ["token"]
-    assert ocr_server.requests == []
-    assert "log_pre_api_call" not in recorder.names
-    if failure == "ordinary":
-        assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
-        assert isinstance(caught.value.__context__, RuntimeError)
-        assert caught.value.__context__.__cause__ is original[failure]
-    else:
-        assert caught.value is original[failure]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-@pytest.mark.parametrize(
-    "override, expected_key",
-    [
-        ({}, "credential-key"),
-        ({"api_key": "explicit-key"}, "explicit-key"),
-        ({"api_key": None}, "environment-key"),
-    ],
-    ids=["inherit", "explicit", "explicit-none"],
-)
-async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments(
-    ocr_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-    asynchronous: bool,
-    override: dict[str, object],
-    expected_key: str,
-) -> None:
-    from litellm.models.credentials import CredentialItem
-
-    pages: Final = [0]
-    opaque: Final = object()
-    monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
-    monkeypatch.setattr(
-        litellm,
-        "credential_list",
-        [
-            CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}),
-            CredentialItem(
-                credential_name="ocr-test",
-                credential_info={},
-                credential_values={
-                    "api_key": "credential-key",
-                    "api_base": ocr_server.base_url,
-                    "pages": pages,
-                    "opaque": opaque,
-                },
-            ),
-            CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}),
-        ],
-    )
-
-    class Observer(RecordingLogger):
-        def log_pre_api_call(self, model, messages, kwargs):
-            super().log_pre_api_call(model, messages, kwargs)
-            pages.append(2)
-
-    arguments: Final = {
-        "model": "mistral/mistral-ocr-latest",
-        "document": OCR_DOCUMENT,
-        "litellm_credential_name": "ocr-test",
-        "callbacks": [Observer()],
-        **override,
-    }
-    response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
-    assert response.pages[0].markdown == "native OCR response"
-    assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}"
-    assert ocr_server.requests[0].body["pages"] == [0, 2]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_native_file_preparation_preserves_reader_exception(
-    ocr_server: RecordingServer, asynchronous: bool
-) -> None:
-    ocr_server.expected_requests = 0
-    failure: Final = RuntimeError("reader failed")
-
-    class Reader:
-        def read(self) -> bytes:
-            raise failure
-
-    document: Final = {"type": "file", "file": Reader()}
-    with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
-        await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr(
-            ocr_server, document=document
-        )
-    assert caught.value.__context__ is failure
-
-
-COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
-FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024
-
-
-class IntReader:
-    def read(self) -> int:
-        return 1
-
-
-def oversized_file(tmp_path: Path) -> Path:
-    path: Final = tmp_path / "large.pdf"
-    with path.open("wb") as stream:
-        stream.truncate(FILE_SIZE_LIMIT + 1)
-    return path
-
-
-def empty_token() -> str:
-    return ""
-
-
-def unused_token() -> str:
-    raise AssertionError("the token provider must not run")
-
-
-@dataclass(frozen=True, slots=True)
-class PublicFailure:
-    arguments: Callable[[Path], dict[str, object]]
-    error: type[Exception]
-    match: str
-    provider_requests: int = 0
-    response: ResponseSpec | None = None
-    cause: type[BaseException] | None = None
-
-
-PUBLIC_FAILURES: Final = {
-    "unknown-req-format": PublicFailure(
-        lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`"
-    ),
-    "empty-file": PublicFailure(
-        lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty"
-    ),
-    "oversized-file": PublicFailure(
-        lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}},
-        litellm.BadRequestError,
-        "exceeds the size limit",
-    ),
-    "missing-file": PublicFailure(
-        lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}},
-        litellm.APIConnectionError,
-        "File not found",
-        cause=FileNotFoundError,
-    ),
-    "reader-returns-non-bytes": PublicFailure(
-        lambda _: {"document": {"type": "file", "file": IntReader()}},
-        litellm.APIConnectionError,
-        "bytes or str",
-        cause=TypeError,
-    ),
-    "cohere-non-image": PublicFailure(
-        lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`"
-    ),
-    "cohere-unknown-format": PublicFailure(
-        lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"},
-        litellm.BadRequestError,
-        "output_format",
-    ),
-    "azure-missing-api-base": PublicFailure(
-        lambda _: {
-            "model": "azure_ai/mistral-ocr-latest",
-            "api_key": None,
-            "api_base": None,
-            "azure_ad_token_provider": unused_token,
-        },
-        litellm.APIConnectionError,
-        "Missing Azure AI API Base",
-    ),
-    "azure-empty-token": PublicFailure(
-        lambda _: {
-            "model": "azure_ai/mistral-ocr-latest",
-            "api_key": None,
-            "azure_ad_token": "static-token",
-            "azure_ad_token_provider": empty_token,
-        },
-        litellm.APIConnectionError,
-        "Missing Azure AI credentials",
-    ),
-    "upstream-500": PublicFailure(
-        lambda _: {},
-        litellm.InternalServerError,
-        "provider unavailable",
-        provider_requests=1,
-        response=ResponseSpec(body={"message": "provider unavailable"}, status=500),
-    ),
-    "invalid-provider-response": PublicFailure(
-        lambda _: {},
-        litellm.APIConnectionError,
-        "pages",
-        provider_requests=1,
-        response=ResponseSpec(body={"pages": "invalid"}),
-    ),
-    "response-over-limit": PublicFailure(
-        lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1},
-        litellm.APIConnectionError,
-        "OCR response exceeds the size limit",
-        provider_requests=1,
-    ),
-    "timeout": PublicFailure(
-        lambda _: {"timeout": 0.01},
-        litellm.Timeout,
-        "",
-        provider_requests=1,
-        response=ResponseSpec(body=OCR_RESPONSE, delay=0.2),
-    ),
-}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys())
-async def test_native_failures_raise_the_public_exception_class(
-    ocr_server: RecordingServer,
-    isolated_azure_auth: None,
-    tmp_path: Path,
-    asynchronous: bool,
-    failure: PublicFailure,
-) -> None:
-    ocr_server.expected_requests = failure.provider_requests
-    if failure.response is not None:
-        ocr_server.enqueue(failure.response)
-
-    with pytest.raises(failure.error, match=failure.match) as caught:
-        await call_native(ocr_server, asynchronous, **failure.arguments(tmp_path))
-
-    assert len(ocr_server.requests) == failure.provider_requests
-    if failure.cause is not None:
-        assert isinstance(caught.value.__context__, failure.cause)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
-    "name,value",
-    [
-        ("ssl_verify", object()),
-        ("ssl_certificate", 1),
-        ("ssl_certificate", ""),
-        ("user_url_allowed_hosts", ["example.test", 1]),
-    ],
-)
-async def test_native_settings_fail_before_provider_io(
-    ocr_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-    asynchronous: bool,
-    name: str,
-    value: object,
-) -> None:
-    ocr_server.expected_requests = 0
-    monkeypatch.setattr(litellm, name, value)
-    with pytest.raises(ValueError, match=r"http_settings|url_policy"):
-        await call_native(ocr_server, asynchronous, num_retries=0)
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
-    "name,value",
-    [
-        ("vertex_project", 1),
-        ("vertex_location", ["region"]),
-    ],
-)
-async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls(
-    ocr_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-    asynchronous: bool,
-    name: str,
-    value: object,
-) -> None:
-    ocr_server.expected_requests = 1
-    monkeypatch.setattr(litellm, name, value)
-    await call_native(ocr_server, asynchronous, num_retries=0)
-    with pytest.raises(ValueError, match=r"provider_defaults"):
-        await call_native(ocr_server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0)
-    assert len(ocr_server.requests) == 1
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_ssl_context_is_terminal_configuration(ocr_server: RecordingServer, asynchronous: bool) -> None:
-    import ssl
-
-    ocr_server.expected_requests = 0
-    context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-    with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"):
-        await call_native(ocr_server, asynchronous, ssl_verify=context, num_retries=0)
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_settings_preserve_protocol_failures(
-    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
-) -> None:
-    ocr_server.expected_requests = 0
-    failure: Final = LookupError("settings truth test failed")
-    cause: Final = RuntimeError("settings cause")
-
-    class RaisesBool:
-        def __bool__(self) -> bool:
-            raise failure from cause
-
-    monkeypatch.setattr(litellm, "force_ipv4", RaisesBool())
-    with pytest.raises(LookupError) as caught:
-        await call_native(ocr_server, asynchronous, num_retries=0)
-    assert caught.value is failure
-    assert caught.value.__cause__ is cause
-    assert caught.value.__traceback__ is not None
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_settings_observe_mutation_between_calls(
-    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
-) -> None:
-    monkeypatch.setattr(litellm, "force_ipv4", "yes")
-    monkeypatch.setattr(litellm, "http2", 1)
-    monkeypatch.setattr(litellm, "vertex_project", [])
-    monkeypatch.setattr(litellm, "vertex_location", 0)
-    monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.")
-    response: Final = await call_native(ocr_server, asynchronous, num_retries=0)
-    assert response.pages[0].markdown == "native OCR response"
-    assert_served_natively(ocr_server, response)
-    monkeypatch.setattr(litellm, "ssl_certificate", 1)
-    with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"):
-        await call_native(ocr_server, asynchronous, num_retries=0)
-    assert len(ocr_server.requests) == 1
-
-
-@pytest.mark.parametrize("required", [False, True])
-@pytest.mark.parametrize("failure", ["invalid", "live", "schema"])
-def test_native_projection_errors_never_select_python(
-    ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str
-) -> None:
-    import dataclasses
-    import ssl
-
-    from litellm.rust_bridge import runtime
-    from litellm.rust_bridge.host import settings
-    from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
-    from litellm.rust_bridge.configuration import Rollout
-    from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR
-    from litellm.rust_bridge.public_call import NativeCall
-
-    ocr_server.expected_requests = 0
-    snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1)
-    if failure == "schema":
-        monkeypatch.setattr(settings, "http_settings", lambda: snapshot)
-    else:
-        monkeypatch.setattr(
-            litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object()
-        )
-    request: Final = NativeCall(
-        args=(),
-        kwargs={},
-        base={
-            "model": "mistral/mistral-ocr-latest",
-            "document": OCR_DOCUMENT,
-            "api_key": "test-key",
-            "api_base": ocr_server.base_url,
-            "timeout": None,
-            "custom_llm_provider": "mistral",
-            "extra_headers": None,
-        },
-    )
-
-    def python_fallback() -> NoReturn:
-        pytest.fail("projection failures must not select Python")
-
-    with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"):
-        runtime.run(
-            RouteContext(Route.OCR, provider="mistral"),
-            binding=NATIVE_OCR,
-            native=lambda native: native(request),
-            python=python_fallback,
-            rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),),
-        )
-    assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"])
-async def test_native_client_certificate_is_validated_before_io(
-    ocr_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-    asynchronous: bool,
-    present: bool,
-) -> None:
-    ocr_server.expected_requests = 0
-    certificate: Final = tmp_path / "client.pem"
-    if present:
-        certificate.write_text("invalid certificate")
-    monkeypatch.setattr(litellm, "ssl_certificate", str(certificate))
-    with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught:
-        await call_native(ocr_server, asynchronous, num_retries=0)
-    assert str(certificate) not in str(caught.value)
-    assert ocr_server.requests == []

tests/test_litellm_rust/ocr/test_settings.py

@@ -0,0 +1,198 @@
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+import litellm
+from tests.test_litellm_rust.ocr.data import (
+    OCR_DOCUMENT,
+    call,
+)
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize(
+    "override, expected_key",
+    [
+        ({}, "credential-key"),
+        ({"api_key": "explicit-key"}, "explicit-key"),
+        ({"api_key": None}, "environment-key"),
+    ],
+    ids=["inherit", "explicit", "explicit-none"],
+)
+async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments(
+    server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    asynchronous: bool,
+    override: dict[str, object],
+    expected_key: str,
+) -> None:
+    from litellm.models.credentials import CredentialItem
+
+    pages: Final = [0]
+    opaque: Final = object()
+    monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
+    monkeypatch.setattr(
+        litellm,
+        "credential_list",
+        [
+            CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}),
+            CredentialItem(
+                credential_name="ocr-test",
+                credential_info={},
+                credential_values={
+                    "api_key": "credential-key",
+                    "api_base": server.base_url,
+                    "pages": pages,
+                    "opaque": opaque,
+                },
+            ),
+            CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}),
+        ],
+    )
+
+    class Observer(RecordingLogger):
+        def log_pre_api_call(self, model, messages, kwargs):
+            super().log_pre_api_call(model, messages, kwargs)
+            pages.append(2)
+
+    arguments: Final = {
+        "model": "mistral/mistral-ocr-latest",
+        "document": OCR_DOCUMENT,
+        "litellm_credential_name": "ocr-test",
+        "callbacks": [Observer()],
+        **override,
+    }
+    response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
+    assert response.pages[0].markdown == "native OCR response"
+    assert server.requests[0].headers["authorization"] == f"Bearer {expected_key}"
+    assert server.requests[0].body["pages"] == [0, 2]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    "name,value",
+    [
+        ("ssl_verify", object()),
+        ("ssl_certificate", 1),
+        ("ssl_certificate", ""),
+        ("user_url_allowed_hosts", ["example.test", 1]),
+    ],
+)
+async def test_native_settings_fail_before_provider_io(
+    server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    asynchronous: bool,
+    name: str,
+    value: object,
+) -> None:
+    server.expected_requests = 0
+    monkeypatch.setattr(litellm, name, value)
+    with pytest.raises(ValueError, match=r"http_settings|url_policy"):
+        await call(server, asynchronous, num_retries=0)
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    "name,value",
+    [
+        ("vertex_project", 1),
+        ("vertex_location", ["region"]),
+    ],
+)
+async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls(
+    server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    asynchronous: bool,
+    name: str,
+    value: object,
+) -> None:
+    server.expected_requests = 1
+    monkeypatch.setattr(litellm, name, value)
+    await call(server, asynchronous, num_retries=0)
+    with pytest.raises(ValueError, match=r"provider_defaults"):
+        await call(server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0)
+    assert len(server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_ssl_context_is_terminal_configuration(server: RecordingServer, asynchronous: bool) -> None:
+    import ssl
+
+    server.expected_requests = 0
+    context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"):
+        await call(server, asynchronous, ssl_verify=context, num_retries=0)
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_settings_preserve_protocol_failures(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+    server.expected_requests = 0
+    failure: Final = LookupError("settings truth test failed")
+    cause: Final = RuntimeError("settings cause")
+
+    class RaisesBool:
+        def __bool__(self) -> bool:
+            raise failure from cause
+
+    monkeypatch.setattr(litellm, "force_ipv4", RaisesBool())
+    with pytest.raises(LookupError) as caught:
+        await call(server, asynchronous, num_retries=0)
+    assert caught.value is failure
+    assert caught.value.__cause__ is cause
+    assert caught.value.__traceback__ is not None
+    assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_settings_observe_mutation_between_calls(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+    monkeypatch.setattr(litellm, "force_ipv4", "yes")
+    monkeypatch.setattr(litellm, "http2", 1)
+    monkeypatch.setattr(litellm, "vertex_project", [])
+    monkeypatch.setattr(litellm, "vertex_location", 0)
+    monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.")
+    response: Final = await call(server, asynchronous, num_retries=0)
+    assert response.pages[0].markdown == "native OCR response"
+    assert_served_natively(server)
+    monkeypatch.setattr(litellm, "ssl_certificate", 1)
+    with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"):
+        await call(server, asynchronous, num_retries=0)
+    assert len(server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"])
+async def test_native_client_certificate_is_validated_before_io(
+    server: RecordingServer,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    asynchronous: bool,
+    present: bool,
+) -> None:
+    server.expected_requests = 0
+    certificate: Final = tmp_path / "client.pem"
+    if present:
+        certificate.write_text("invalid certificate")
+    monkeypatch.setattr(litellm, "ssl_certificate", str(certificate))
+    with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught:
+        await call(server, asynchronous, num_retries=0)
+    assert str(certificate) not in str(caught.value)
+    assert server.requests == []

tests/test_litellm_rust/responses/conftest.py

@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+    recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
+    return recording_server

tests/test_litellm_rust/responses/data.py

@@ -0,0 +1,75 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+RESPONSES_MODEL: Final = "openai/gpt-6.1-sol"
+RESPONSES_TEXT: Final = "Hello from native Responses"
+RESPONSES_RESPONSE: Final = {
+    "id": "resp_native",
+    "object": "response",
+    "created_at": 1,
+    "status": "completed",
+    "model": RESPONSES_MODEL.removeprefix("openai/"),
+    "output": [
+        {
+            "id": "msg_native",
+            "type": "message",
+            "role": "assistant",
+            "status": "completed",
+            "content": [{"type": "output_text", "text": RESPONSES_TEXT, "annotations": []}],
+        }
+    ],
+    "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9},
+}
+RESPONSES_EVENTS: Final = (
+    (
+        "response.created",
+        {
+            "type": "response.created",
+            "sequence_number": 0,
+            "response": {**RESPONSES_RESPONSE, "status": "in_progress", "output": []},
+        },
+    ),
+    (
+        "response.output_text.delta",
+        {
+            "type": "response.output_text.delta",
+            "sequence_number": 1,
+            "item_id": "msg_native",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": RESPONSES_TEXT,
+        },
+    ),
+    ("response.completed", {"type": "response.completed", "sequence_number": 2, "response": RESPONSES_RESPONSE}),
+)
+
+
+def _error(kind: str) -> dict[str, object]:
+    return {"error": {"type": kind, "message": kind, "param": None, "code": None}}
+
+
+REQUEST: Final = {"model": RESPONSES_MODEL, "input": "Hello"}
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_exceeded")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+    return await call_natively(
+        Route.RESPONSES,
+        litellm.responses,
+        litellm.aresponses,
+        REQUEST,
+        server,
+        asynchronous=asynchronous,
+        **kwargs,
+    )
+
+
+SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook")
+STREAMING: Final = gap("native Responses streaming declines to Python", raises=RuntimeError)
+BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks")

tests/test_litellm_rust/responses/test_body_hooks.py

@@ -0,0 +1,114 @@
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Final
+
+import pytest
+from pydantic import BaseModel
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from tests.test_litellm_rust.responses.data import (
+    BODY_HOOKS,
+    REQUEST,
+    RESPONSES_EVENTS,
+    RESPONSES_RESPONSE,
+    STREAMING,
+    call,
+)
+from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume, response_id
+from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_body_streaming_deployment_hook_rewrites_chunks_before_delivery(server: RecordingServer) -> None:
+    seen: Final[list[str]] = []
+
+    class Rewrite(CustomLogger):
+        async def async_post_call_streaming_deployment_hook(
+            self, request_data: dict[str, object], response_chunk: BaseModel, call_type: object
+        ) -> BaseModel:
+            seen.append(getattr(response_chunk, "type", ""))
+            if getattr(response_chunk, "type", None) != "response.output_text.delta":
+                return response_chunk
+            return response_chunk.model_copy(update={"delta": "rewritten"})
+
+    server.enqueue(STREAM)
+    with rebound(litellm, "callbacks", [Rewrite()]):
+        stream: Final = await call(server, stream=True)
+        chunks: Final = await consume(stream)
+
+    assert_served_natively(server)
+    assert seen == [event for event, _ in RESPONSES_EVENTS]
+    assert [getattr(chunk, "delta", None) for chunk in chunks] == [None, "rewritten", None]
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None:
+    loop: Final = PlannedAgenticLoop()
+    server.enqueue(ResponseSpec(body=RESPONSES_RESPONSE))
+    server.default_response = ResponseSpec(body={**RESPONSES_RESPONSE, "id": "followup"})
+
+    with rebound(litellm, "callbacks", [loop]):
+        response: Final = await call(server)
+        await drain_logging()
+
+    assert_served_natively(server, count=2)
+    assert isinstance(server.requests[1].body, dict)
+    assert server.requests[1].body["input"] == [{"role": "user", "content": "followup"}]
+    assert response_id(response) == "agentic-replaced"
+    assert agentic_hooks(loop) == (
+        "async_should_run_agentic_loop",
+        "async_build_agentic_loop_plan",
+        "async_should_run_agentic_loop",
+        "async_post_agentic_loop_response_hook",
+        "async_agentic_loop_cleanup_hook",
+    )
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None:
+    loop: Final = LegacyAgenticLoop()
+
+    with rebound(litellm, "callbacks", [loop]):
+        response: Final = await call(server)
+
+    assert_served_natively(server)
+    assert response_id(response) == "agentic-replaced"
+    assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop")
+
+
+@pytest.mark.asyncio
+async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None:
+    observed: Final[list[object]] = []
+    logging_obj: Final = Logging(
+        model=REQUEST["model"],
+        messages=[],
+        stream=False,
+        start_time=datetime.now(),
+        call_type="aresponses",
+        litellm_call_id=CALL_ID,
+        function_id=CALL_ID,
+    )
+    logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={})
+
+    await call(server, litellm_logging_obj=logging_obj)
+
+    assert_served_natively(server)
+    assert observed == []

tests/test_litellm_rust/responses/test_callbacks.py

@@ -0,0 +1,408 @@
+import asyncio
+from collections.abc import Iterator
+from contextvars import ContextVar
+from datetime import datetime
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.argilla import ArgillaLogger
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.integrations.argilla import ArgillaItem
+from litellm.types.utils import StandardLoggingPayload
+from tests.test_litellm_rust.responses.data import (
+    BAD_REQUEST,
+    RATE_LIMITED,
+    RESPONSES_RESPONSE,
+    SYNC_DEPLOYMENT_HOOK,
+    call,
+)
+from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id, with_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers")
+_OBJECT: Final = TypeAdapter(dict[str, object])
+_HEADERS: Final = TypeAdapter(dict[str, str])
+
+
+def success_signal(asynchronous: bool) -> str:
+    return "async_log_success_event" if asynchronous else "log_success_event"
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "hook,asynchronous,expected",
+    (
+        pytest.param("log_pre_api_call", True, 1, id="async-log_pre_api_call"),
+        pytest.param("log_post_api_call", True, 1, id="async-log_post_api_call"),
+        pytest.param("async_log_success_event", True, 1, id="async-async_log_success_event"),
+        pytest.param("async_logging_hook", True, 1, id="async-async_logging_hook"),
+        pytest.param("redact_standard_logging_payload_from_model_call_details", True, 1, id="async-redact"),
+        pytest.param("async_pre_call_deployment_hook", True, 1, id="async-async_pre_call_deployment_hook"),
+        pytest.param(
+            "async_post_call_success_deployment_hook", True, 1, id="async-async_post_call_success_deployment_hook"
+        ),
+        pytest.param("log_success_event", True, 0, id="async-log_success_event"),
+        pytest.param("logging_hook", True, 0, id="async-logging_hook"),
+        pytest.param("log_pre_api_call", False, 1, id="sync-log_pre_api_call"),
+        pytest.param("log_post_api_call", False, 1, id="sync-log_post_api_call"),
+        pytest.param("log_success_event", False, 1, id="sync-log_success_event"),
+        pytest.param("logging_hook", False, 1, id="sync-logging_hook"),
+        pytest.param(
+            "async_pre_call_deployment_hook",
+            False,
+            1,
+            id="sync-async_pre_call_deployment_hook",
+            marks=SYNC_DEPLOYMENT_HOOK,
+        ),
+        pytest.param(
+            "async_post_call_success_deployment_hook", False, 0, id="sync-async_post_call_success_deployment_hook"
+        ),
+    ),
+)
+async def test_success_hook_fires_the_expected_number_of_times(
+    server: RecordingServer, hook: str, asynchronous: bool, expected: int
+) -> None:
+    recorder: Final = HookRecorder()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        await call(server, asynchronous=asynchronous)
+        await recorder.wait_for_async(success_signal(asynchronous))
+
+    assert_served_natively(server)
+    assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "hook,asynchronous,expected",
+    (
+        pytest.param("async_log_failure_event", True, 1, id="async-async_log_failure_event"),
+        pytest.param(
+            "async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook"
+        ),
+        pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"),
+        pytest.param("log_failure_event", True, 0, id="async-log_failure_event", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+        pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"),
+    ),
+)
+async def test_failure_hook_fires_the_expected_number_of_times(
+    server: RecordingServer, hook: str, asynchronous: bool, expected: int
+) -> None:
+    server.default_response = ResponseSpec(body=RATE_LIMITED, status=429)
+    recorder: Final = HookRecorder()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        with pytest.raises(litellm.RateLimitError):
+            await call(server, asynchronous=asynchronous, metadata={"model_group_size": 1})
+        await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
+
+    assert_served_natively(server)
+    assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None:
+    server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+    recorder: Final = RecordingLogger()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        with pytest.raises(litellm.BadRequestError) as raised:
+            await call(server)
+        events: Final = await recorder.wait_for_async("async_log_failure_event")
+
+    assert_served_natively(server)
+    assert isinstance(events[0].kwargs, dict)
+    assert events[0].kwargs["exception"] is raised.value
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    await call(server, callbacks=[recorder])
+
+    assert_served_natively(server)
+    pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
+    assert isinstance(pre_call, dict)
+    assert request_body(pre_call) == server.requests[0].body
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
+    class Edit(CustomLogger):
+        def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+            request_body(kwargs)["temperature"] = 0.25
+
+    await call(server, callbacks=[Edit()])
+
+    assert_served_natively(server)
+    assert isinstance(server.requests[0].body, dict)
+    assert server.requests[0].body["temperature"] == 0.25
+
+
+@pytest.mark.asyncio
+async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None:
+    class Rewrite(HookRecorder):
+        async def async_logging_hook(
+            self, kwargs: dict[str, object], result: object, call_type: str
+        ) -> tuple[dict[str, object], object]:
+            return {**kwargs, "logging_hook_ran": True}, result
+
+        def redact_standard_logging_payload_from_model_call_details(
+            self, model_call_details: dict[str, object]
+        ) -> dict[str, object]:
+            return {**model_call_details, "redacted": True}
+
+    recorder: Final = Rewrite()
+
+    with rebound(litellm, "callbacks", [recorder]):
+        await call(server)
+        success: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert isinstance(success[0].kwargs, dict)
+    assert success[0].kwargs["logging_hook_ran"] is True
+    assert success[0].kwargs["redacted"] is True
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order(
+    server: RecordingServer,
+) -> None:
+    class Rewrite(CustomLogger):
+        def __init__(self, marker: str) -> None:
+            super().__init__()
+            self.marker = marker
+
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            headers: Final = _HEADERS.validate_python(kwargs["extra_headers"])
+            return {
+                **kwargs,
+                "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker},
+            }
+
+    with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]):
+        await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"})
+
+    assert_served_natively(server)
+    assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second"
+    assert server.requests[0].headers["x-kept"] == "unchanged"
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None:
+    rejection: Final = litellm.BadRequestError(message="deployment rejected", model="test-model", llm_provider="openai")
+    recorder: Final = RecordingLogger()
+
+    class Reject(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            raise rejection
+
+    with rebound(litellm, "callbacks", [Reject(), recorder]):
+        with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised:
+            await call(server)
+        failures: Final = await recorder.wait_for_async("async_log_failure_event")
+
+    assert raised.value is rejection
+    assert server.requests == []
+    assert len(failures) == 1
+    assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value
+    assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names)
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
+    class DropKeywords(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            return {name: value for name, value in kwargs.items() if name not in ("instructions", "extra_headers")}
+
+    with rebound(litellm, "callbacks", [DropKeywords()]):
+        await call(
+            server,
+            instructions="remove this instruction",
+            extra_headers={"x-deleted": "remove this header"},
+        )
+
+    assert_served_natively(server)
+    sent: Final = server.requests[0]
+    assert isinstance(sent.body, dict)
+    assert "instructions" not in sent.body
+    assert "x-deleted" not in sent.headers
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None:
+    recorder: Final = RecordingLogger()
+
+    class Replace(CustomLogger):
+        async def async_post_call_success_deployment_hook(
+            self, request_data: dict[str, object], response: object, call_type: object
+        ) -> object:
+            return with_id(response, "replaced")
+
+    with rebound(litellm, "callbacks", [Replace(), recorder]):
+        response: Final = await call(server)
+        success: Final = await recorder.wait_for_async("async_log_success_event")
+
+    assert_served_natively(server)
+    assert response_id(response) == "replaced"
+    assert len(success) == 1
+    assert response_id(success[0].response) == response_id(response)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None:
+    context: Final = ContextVar("responses-callback-context", default="outside")
+    first_entered: Final = asyncio.Event()
+    second_entered: Final = asyncio.Event()
+    server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_first"}))
+    server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_second"}))
+
+    class Observe(RecordingLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: object
+        ) -> dict[str, object]:
+            arrived, other = (
+                (first_entered, second_entered)
+                if kwargs["litellm_call_id"] == "first"
+                else (second_entered, first_entered)
+            )
+            arrived.set()
+            await asyncio.wait_for(other.wait(), timeout=5)
+            return kwargs
+
+        def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+            self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()})
+
+        async def async_log_success_event(
+            self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+        ) -> None:
+            self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj)
+
+    recorder: Final = Observe()
+
+    async def invoke(marker: str) -> object:
+        token: Final = context.set(marker)
+        try:
+            return await call(server, litellm_call_id=marker, metadata={"marker": marker})
+        finally:
+            context.reset(token)
+
+    with rebound(litellm, "callbacks", [recorder]):
+        responses: Final = await asyncio.gather(invoke("first"), invoke("second"))
+        success: Final = await recorder.wait_for_async("async_log_success_event", count=2)
+
+    assert_served_natively(server, count=2)
+    expected: Final = {"first": responses[0], "second": responses[1]}
+    pre_call: Final = recorder.wait_for("log_pre_api_call", count=2)
+    assert len(pre_call) == len(success) == 2
+    assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys()
+    for event in (*pre_call, *success):
+        details: Final = _OBJECT.validate_python(event.kwargs)
+        call_id: Final = str(details["litellm_call_id"])
+        params: Final = _OBJECT.validate_python(details["litellm_params"])
+        assert details["caller_context"] == call_id
+        assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id
+    for event in success:
+        logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"])
+        assert response_id(event.response) == response_id(expected[logged_call_id])
+    assert context.get() == "outside"
+
+
+@pytest.mark.asyncio
+async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None:
+    class Broken(CustomLogger):
+        def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+            raise RuntimeError("logger exploded")
+
+    response: Final = await call(server, callbacks=[Broken()])
+
+    assert_served_natively(server)
+    assert response_id(response) == RESPONSES_RESPONSE["id"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "event,asynchronous,dynamic,expected",
+    (
+        pytest.param("input", False, False, 1, id="global-input"),
+        pytest.param("success", True, False, 1, id="global-async-success"),
+        pytest.param("success", True, True, 1, id="request-async-success"),
+        pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+        pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+        pytest.param("failure", True, False, 1, id="global-async-failure"),
+        pytest.param("failure", False, False, 1, id="global-sync-failure"),
+        pytest.param("failure", False, True, 1, id="request-sync-failure"),
+    ),
+)
+async def test_function_callback_fires_the_expected_number_of_times(
+    server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int
+) -> None:
+    recorder: Final = RecordingLogger()
+
+    def callback(kwargs: dict[str, object], *_: object) -> None:
+        recorder._record(event, kwargs)
+
+    async def async_callback(kwargs: dict[str, object], *_: object) -> None:
+        recorder._record(event, kwargs)
+
+    registered: Final = async_callback if asynchronous else callback
+    global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback"
+    request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {}
+    if event == "failure":
+        server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+
+    with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]):
+        try:
+            await call(server, **request_callbacks)
+        except litellm.BadRequestError:
+            assert event == "failure"
+        await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event")
+
+    assert_served_natively(server)
+    assert recorder.names.count(event) == expected
+
+
+@pytest.mark.asyncio
+async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None:
+    datasets: Final = RecordingLogger()
+
+    class Drop(CustomLogger):
+        async def async_dataset_hook(
+            self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None
+        ) -> None:
+            datasets._record("async_dataset_hook", response=logged_data)
+
+    server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]}))
+    with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}):
+        sink: Final = ArgillaLogger(
+            argilla_api_key="test-key",
+            argilla_base_url=server.base_url,
+            argilla_dataset_name="callback-dataset",
+            flush_interval=3600,
+        )
+        with rebound(litellm, "callbacks", [sink, Drop()]):
+            await call(server)
+            await datasets.wait_for_async("async_dataset_hook")
+            await sink.flush_queue()
+
+    assert [request.method for request in server.requests] == ["GET", "POST"]

tests/test_litellm_rust/responses/test_failures.py

@@ -0,0 +1,83 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm import RateLimitError
+from tests.test_litellm_rust.responses.data import RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+def failure_signal(asynchronous: bool) -> str:
+    return "async_log_failure_event" if asynchronous else "log_failure_event"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+    monkeypatch.setattr(litellm, "max_budget", 1.0)
+    monkeypatch.setattr(litellm, "_current_cost", 2.0)
+    monkeypatch.setattr(litellm, "ssl_verify", object())
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+
+    with pytest.raises(litellm.BudgetExceededError) as caught:
+        await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+    failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not server.requests
+    assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_provider_failure_is_terminal_and_shared_with_callbacks(
+    server: RecordingServer, asynchronous: bool
+) -> None:
+    recorder: Final = RecordingLogger()
+    server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429))
+
+    with pytest.raises(RateLimitError) as caught:
+        await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+    assert caught.value.status_code == 429
+    assert len(server.requests) == 1
+    failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "options",
+    (
+        {"stream": True},
+        {"extra_body": {"provider_option": True}},
+        {"mock_response": "mock"},
+        {"num_retries": 1},
+        {"use_chat_completions_api": True},
+        {"model_list": []},
+    ),
+)
+async def test_unsupported_requests_fail_before_callbacks_and_provider_io(
+    server: RecordingServer, options: dict[str, object]
+) -> None:
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+
+    with pytest.raises(RuntimeError, match="declined the request"):
+        await call(server, callbacks=[recorder], **options)
+
+    assert not server.requests
+    assert not recorder.events

tests/test_litellm_rust/responses/test_lifecycle.py

@@ -0,0 +1,46 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge import _native
+from litellm.rust_bridge.public_call import NativeCall
+from tests.test_litellm_rust.responses.data import REQUEST
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+async def test_unstarted_call_has_no_provider_or_callback_effects(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+    monkeypatch.setattr(litellm, "ssl_verify", object())
+    monkeypatch.setattr(litellm, "max_budget", 1.0)
+    monkeypatch.setattr(litellm, "_current_cost", 2.0)
+    kwargs: Final = {
+        **REQUEST,
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        "callbacks": [recorder],
+    }
+
+    base: Final = {
+        **REQUEST,
+        "stream": None,
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        "custom_llm_provider": "openai",
+        "extra_headers": None,
+    }
+
+    pending: Final = _native.aresponses(NativeCall(args=(), kwargs=kwargs, base=base))
+    assert asyncio.iscoroutine(pending)
+    pending.close()
+
+    assert not server.requests
+    assert not recorder.events

tests/test_litellm_rust/responses/test_requests.py

@@ -0,0 +1,56 @@
+import asyncio
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.rust_bridge.catalog import Route
+from litellm.types.llms.openai import ResponsesAPIResponse
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_TEXT, call
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None:
+    arguments: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16)
+    supplied: Final = {"api_key": "test-key", "api_base": server.base_url}
+
+    with native_route(Route.RESPONSES):
+        result: Final = (
+            await litellm.aresponses(*arguments, **supplied)
+            if asynchronous
+            else await asyncio.to_thread(litellm.responses, *arguments, **supplied)
+        )
+
+    assert isinstance(result, ResponsesAPIResponse)
+    assert result.output_text == RESPONSES_TEXT
+    assert_served_natively(server)
+    body: Final = _OBJECT.validate_python(server.requests[0].body)
+    assert body["instructions"] == arguments[3]
+    assert body["max_output_tokens"] == arguments[4]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("encoded", (False, True))
+async def test_continuation_ids_are_decoded_before_they_reach_the_wire(
+    server: RecordingServer, asynchronous: bool, encoded: bool
+) -> None:
+    original: Final = "resp_upstream"
+    previous: Final = (
+        ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original)
+        if encoded
+        else original
+    )
+
+    await call(server, asynchronous=asynchronous, previous_response_id=previous)
+
+    assert _OBJECT.validate_python(server.requests[0].body)["previous_response_id"] == original

tests/test_litellm_rust/responses/test_settings.py

@@ -0,0 +1,83 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.models.credentials import CredentialItem
+from litellm.types.llms.openai import ResponsesAPIResponse
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.responses.data import call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("from_credentials", (False, True))
+async def test_resource_setup_uses_deployment_hook_arguments(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool
+) -> None:
+    invalid_settings: Final = {"ssl_verify": object()}
+    credential: Final = CredentialItem(
+        credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
+    )
+    monkeypatch.setattr(litellm, "credential_list", [credential])
+
+    class Prepare(CustomLogger):
+        async def async_pre_call_deployment_hook(
+            self, kwargs: dict[str, object], call_type: CallTypes | None
+        ) -> dict[str, object]:
+            return {
+                **kwargs,
+                **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
+            }
+
+    litellm.callbacks.append(Prepare())
+    recorder: Final = RecordingLogger()
+    server.expected_requests = 0
+
+    with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
+        await call(server, callbacks=[recorder])
+
+    failure: Final = await recorder.wait_for_async("async_log_failure_event")
+    assert len(failure) == 1
+    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+    assert not server.requests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
+async def test_connection_settings_reach_the_provider(
+    server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str
+) -> None:
+    key: Final = "selected-key"
+    explicit: Final = source in ("explicit", "base_url")
+    monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
+    monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None))
+    monkeypatch.setattr(
+        litellm,
+        "api_base",
+        None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url),
+    )
+    monkeypatch.setenv("OPENAI_API_KEY", key if source == "environment" else "unused")
+    for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE"):
+        monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1")
+
+    result: Final = await call(
+        server,
+        asynchronous=asynchronous,
+        api_key=key if explicit else ("" if source == "empty" else None),
+        api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None),
+        **({"base_url": server.base_url} if source == "base_url" else {}),
+    )
+
+    assert isinstance(result, ResponsesAPIResponse)
+    assert_served_natively(server)
+    assert server.requests[0].headers["authorization"] == f"Bearer {key}"

tests/test_litellm_rust/responses/test_streaming.py

@@ -0,0 +1,58 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.responses.data import RESPONSES_EVENTS, RESPONSES_RESPONSE, STREAMING, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume, response_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+    yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
+    server: RecordingServer,
+) -> None:
+    server.enqueue(STREAM)
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, stream=True, callbacks=[recorder])
+    assert isinstance(stream, AsyncIterator)
+    first: Final = await anext(stream)
+    await drain_logging()
+    assert "async_log_success_event" not in recorder.names
+    rest: Final = [chunk async for chunk in stream]
+
+    assert_served_natively(server)
+    assert [getattr(chunk, "type", None) for chunk in (first, *rest)] == [event for event, _ in RESPONSES_EVENTS]
+    success: Final = await recorder.wait_for_async("async_log_success_event")
+    assert len(success) == 1
+    assert isinstance(success[0].kwargs, dict)
+    assert success[0].kwargs["completion_start_time"] is not None
+    assert response_id(success[0].response) == RESPONSES_RESPONSE["id"]
+    assert "log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None:
+    server.enqueue(STREAM)
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder])
+    chunks: Final = await consume(stream)
+
+    assert_served_natively(server)
+    assert [getattr(chunk, "type", None) for chunk in chunks] == [event for event, _ in RESPONSES_EVENTS]
+    assert len(await recorder.wait_for_async("log_success_event")) == 1

tests/test_litellm_rust/support/body_hooks.py

@@ -0,0 +1,39 @@
+from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.results import response_id, with_id
+
+
+def agentic_hooks(recorder: RecordingLogger) -> tuple[str, ...]:
+    return tuple(name for name in recorder.names if "agentic" in name)
+
+
+class PlannedAgenticLoop(RecordingLogger):
+    async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]:
+        self._record("async_should_run_agentic_loop", response=response)
+        return response_id(response) != "followup", {}
+
+    async def async_build_agentic_loop_plan(self, response: object, **kwargs: object) -> AgenticLoopPlan:
+        self._record("async_build_agentic_loop_plan", response=response)
+        return AgenticLoopPlan(
+            run_agentic_loop=True,
+            request_patch=AgenticLoopRequestPatch.model_validate(
+                {"messages": [{"role": "user", "content": "followup"}]}
+            ),
+        )
+
+    async def async_post_agentic_loop_response_hook(self, response: object, **kwargs: object) -> object:
+        self._record("async_post_agentic_loop_response_hook", response=response)
+        return with_id(response, "agentic-replaced")
+
+    async def async_agentic_loop_cleanup_hook(self, **kwargs: object) -> None:
+        self._record("async_agentic_loop_cleanup_hook")
+
+
+class LegacyAgenticLoop(RecordingLogger):
+    async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]:
+        self._record("async_should_run_agentic_loop", response=response)
+        return True, {}
+
+    async def async_run_agentic_loop(self, response: object, **kwargs: object) -> object:
+        self._record("async_run_agentic_loop", response=response)
+        return with_id(response, "agentic-replaced")

tests/test_litellm_rust/support/callback_recorder.py

@@ -3,6 +3,7 @@ import copy
 import threading
 import time
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Final
 
 from litellm.integrations.custom_logger import CustomLogger
@@ -99,3 +100,63 @@ class RecordingLogger(CustomLogger):
     def logging_hook(self, kwargs, result, call_type):
         self._record("logging_hook", kwargs, result)
         return kwargs, result
+
+
+class HookRecorder(RecordingLogger):
+    def log_post_api_call(
+        self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+    ) -> None:
+        self._record("log_post_api_call", kwargs, response_obj)
+
+    async def async_logging_hook(
+        self, kwargs: dict[str, object], result: object, call_type: str
+    ) -> tuple[dict[str, object], object]:
+        self._record("async_logging_hook", kwargs, result)
+        return kwargs, result
+
+    def redact_standard_logging_payload_from_model_call_details(
+        self, model_call_details: dict[str, object]
+    ) -> dict[str, object]:
+        self._record("redact_standard_logging_payload_from_model_call_details", model_call_details)
+        return model_call_details
+
+    async def async_pre_call_deployment_hook(self, kwargs: dict[str, object], call_type: object) -> dict[str, object]:
+        self._record("async_pre_call_deployment_hook", kwargs)
+        return kwargs
+
+    async def async_post_call_success_deployment_hook(
+        self, request_data: dict[str, object], response: object, call_type: object
+    ) -> object:
+        self._record("async_post_call_success_deployment_hook", request_data, response)
+        return response
+
+    async def async_post_call_failure_deployment_hook(
+        self,
+        request_data: dict[str, object],
+        original_exception: Exception,
+        call_type: object,
+        *,
+        fallback_depth: int | None = None,
+    ) -> None:
+        self._record("async_post_call_failure_deployment_hook", request_data, original_exception)
+
+    async def log_model_group_rate_limit_error(
+        self, exception: Exception, original_model_group: str | None, kwargs: dict[str, object]
+    ) -> None:
+        self._record("log_model_group_rate_limit_error", kwargs, exception)
+
+
+def request_body(kwargs: dict[str, object]) -> dict[str, object]:
+    additional_args = kwargs["additional_args"]
+    assert isinstance(additional_args, dict)
+    body = additional_args["complete_input_dict"]
+    assert isinstance(body, dict)
+    return body
+
+
+def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
+    additional_args = kwargs["additional_args"]
+    assert isinstance(additional_args, dict)
+    headers = additional_args["headers"]
+    assert isinstance(headers, dict)
+    return headers

tests/test_litellm_rust/support/gaps.py

@@ -0,0 +1,5 @@
+import pytest
+
+
+def gap(reason: str, raises: type[Exception] = AssertionError) -> pytest.MarkDecorator:
+    return pytest.mark.xfail(reason=reason, raises=raises, strict=True)

tests/test_litellm_rust/support/recording_server.py

@@ -120,3 +120,9 @@ 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 == []
+
+
+def relaxed(server: RecordingServer) -> Iterator[RecordingServer]:
+    server.expected_requests = None
+    yield server
+    server.responses.clear()

tests/test_litellm_rust/support/requests.py

@@ -1,97 +0,0 @@
-from typing import Final
-
-import litellm
-from litellm.llms.base_llm.ocr.transformation import OCRResponse
-from tests.test_litellm_rust.support.recording_server import RecordingServer
-
-OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
-OCR_MODEL: Final = "mistral/mistral-ocr-latest"
-OCR_RESPONSE: Final = {
-    "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
-    "model": "mistral-ocr-latest",
-    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
-}
-
-MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5"
-MESSAGES: Final = ({"role": "user", "content": "Hello"},)
-MESSAGES_RESPONSE: Final = {
-    "id": "msg_native",
-    "type": "message",
-    "role": "assistant",
-    "model": "claude-sonnet-5",
-    "content": [{"type": "text", "text": "Hello from native Messages"}],
-    "stop_reason": "end_turn",
-    "stop_sequence": None,
-    "usage": {"input_tokens": 5, "output_tokens": 4},
-}
-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": {"output_tokens": 4},
-        },
-    ),
-    ("message_stop", {"type": "message_stop"}),
-)
-
-
-def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
-    return {
-        "model": OCR_MODEL,
-        "document": dict(OCR_DOCUMENT),
-        "api_key": "test-key",
-        "api_base": server.base_url,
-        **kwargs,
-    }
-
-
-def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
-    response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
-    if not isinstance(response, OCRResponse):
-        raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
-    return response
-
-
-async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
-    return await litellm.aocr(**ocr_arguments(server, **kwargs))
-
-
-def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
-    return call_ocr(server, **kwargs)
-
-
-async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
-    return await call_aocr(server, **kwargs)
-
-
-async def call_native(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse:
-    return await call_native_aocr(server, **kwargs) if asynchronous else call_native_ocr(server, **kwargs)
-
-
-def request_body(kwargs: dict[str, object]) -> dict[str, object]:
-    additional_args = kwargs["additional_args"]
-    assert isinstance(additional_args, dict)
-    body = additional_args["complete_input_dict"]
-    assert isinstance(body, dict)
-    return body
-
-
-def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
-    additional_args = kwargs["additional_args"]
-    assert isinstance(additional_args, dict)
-    headers = additional_args["headers"]
-    assert isinstance(headers, dict)
-    return headers

tests/test_litellm_rust/support/results.py

@@ -0,0 +1,26 @@
+from collections.abc import AsyncIterator, Iterator
+
+from pydantic import BaseModel
+
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.types.llms.openai import ResponsesAPIResponse
+
+
+def response_id(response: object) -> object:
+    if isinstance(response, ResponsesAPIResponse):
+        return ResponsesAPIRequestUtils.decode_responses_api_response_id(response.id).get("response_id") or response.id
+    return response.get("id") if isinstance(response, dict) else None
+
+
+def with_id(response: object, new_id: str) -> object:
+    if isinstance(response, BaseModel):
+        return response.model_copy(update={"id": new_id})
+    assert isinstance(response, dict)
+    return {**response, "id": new_id}
+
+
+async def consume(stream: object) -> tuple[object, ...]:
+    if isinstance(stream, AsyncIterator):
+        return tuple([chunk async for chunk in stream])
+    assert isinstance(stream, Iterator)
+    return tuple(stream)

tests/test_litellm_rust/support/routing.py

@@ -0,0 +1,48 @@
+import asyncio
+import copy
+from collections.abc import Awaitable, Callable, Generator, Mapping
+from contextlib import contextmanager
+from typing import Final
+
+from litellm.rust_bridge import catalog
+from litellm.rust_bridge.catalog import Route, RouteRule
+from litellm.rust_bridge.configuration import Rollout
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.served import assert_served_natively
+
+__all__ = ["CALL_ID", "assert_served_natively", "call_natively", "native_route"]
+
+CALL_ID: Final = "callback-contract"
+
+
+@contextmanager
+def native_route(*routes: Route) -> Generator[None]:
+    rules: tuple[RouteRule, ...] = tuple(RouteRule(route, Rollout.RUST_REQUIRED) for route in routes)
+    with rebound(catalog, "RULES", (*rules, *catalog.RULES)):
+        yield
+
+
+async def call_natively(
+    route: Route,
+    create: Callable[..., object],
+    acreate: Callable[..., Awaitable[object]],
+    request: Mapping[str, object],
+    server: RecordingServer,
+    *,
+    asynchronous: bool,
+    **kwargs: object,
+) -> object:
+    arguments: Final = {
+        **copy.deepcopy(dict(request)),
+        "api_key": "test-key",
+        "api_base": server.base_url,
+        "litellm_call_id": CALL_ID,
+        **kwargs,
+    }
+    with native_route(route):
+        if asynchronous:
+            return await acreate(**arguments)
+        return await asyncio.to_thread(create, **arguments)
+
+

tests/test_litellm_rust/support/served.py

@@ -4,8 +4,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
 from tests.test_litellm_rust.support.recording_server import RecordingServer
 
 
-def assert_served_natively(server: RecordingServer, response: object) -> None:
-    assert len(server.requests) == 1
-    additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
-    assert isinstance(additional_headers, dict)
-    assert additional_headers.get("x-litellm-rust") == "true"
+def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None:
+    assert len(server.requests) == (len(responses) or count)
+    if not responses:
+        assert not any(
+            request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests
+        )
+    for response in responses:
+        additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
+        assert isinstance(additional_headers, dict)
+        assert additional_headers.get("x-litellm-rust") == "true"

tests/test_litellm_rust/test_inference.py

@@ -1,403 +0,0 @@
-import asyncio
-from collections.abc import Awaitable, Coroutine, Mapping
-from typing import Final, Literal, TypeAlias
-
-import pytest
-from pydantic import JsonValue, TypeAdapter
-
-import litellm
-from litellm import RateLimitError
-from litellm.chat_completions import dispatch as chat_dispatch
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.models.credentials import CredentialItem
-from litellm.responses.utils import ResponsesAPIRequestUtils
-from litellm.rust_bridge import _native
-from litellm.rust_bridge.public_call import NativeCall
-from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.types.utils import CallTypes, ModelResponse
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE, request_body
-
-pytestmark = pytest.mark.requires_rust_extension
-Route: TypeAlias = Literal["chat", "responses"]
-_OBJECT: Final = TypeAdapter(dict[str, object])
-NativeResult: TypeAlias = (
-    ModelResponse
-    | ResponsesAPIResponse
-    | Coroutine[object, object, ModelResponse]
-    | Coroutine[object, object, ResponsesAPIResponse]
-)
-RESPONSES_MODEL: Final = "openai/gpt-6-sol"
-RESPONSES_RESPONSE: Final[dict[str, JsonValue]] = {
-    "id": "resp_native",
-    "object": "response",
-    "created_at": 1,
-    "model": RESPONSES_MODEL.removeprefix("openai/"),
-    "status": "completed",
-    "output": [
-        {
-            "type": "message",
-            "id": "msg_native",
-            "role": "assistant",
-            "status": "completed",
-            "content": [{"type": "output_text", "text": "native response", "annotations": []}],
-        }
-    ],
-    "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9},
-}
-
-
-@pytest.fixture(params=("chat", "responses"))
-def route(request: pytest.FixtureRequest) -> Route:
-    return TypeAdapter(Route).validate_python(request.param)
-
-
-def native_call(
-    route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]
-) -> NativeResult:
-    server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE)
-    if route == "chat":
-        kwargs: Final = {
-            "model": MESSAGES_MODEL,
-            "messages": list(MESSAGES),
-            "api_key": "test-key",
-            "api_base": server.base_url,
-            "max_tokens": 32,
-            **options,
-        }
-        request: Final = NativeCall(args=(), kwargs=kwargs, base={})
-        return (_native.acompletion if asynchronous else _native.completion)(request)
-    response_kwargs: Final = {
-        "model": RESPONSES_MODEL,
-        "input": "hello",
-        "api_key": "test-key",
-        "api_base": server.base_url,
-        "max_output_tokens": 32,
-        **options,
-    }
-    response_request: Final = NativeCall(
-        args=(),
-        kwargs=response_kwargs,
-        base={
-            "model": RESPONSES_MODEL,
-            "input": "hello",
-            "stream": None,
-            "api_key": "test-key",
-            "api_base": server.base_url,
-            "custom_llm_provider": "openai",
-            "extra_headers": None,
-        },
-    )
-    return (_native.aresponses if asynchronous else _native.responses)(response_request)
-
-
-async def execute(route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]) -> object:
-    if not asynchronous:
-        return await asyncio.to_thread(native_call, route, False, server, options)
-    result: Final = native_call(route, True, server, options)
-    assert isinstance(result, Awaitable)
-    return await result
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_inference_returns_public_models_and_logs_once(
-    route: Route,
-    asynchronous: bool,
-    recording_server: RecordingServer,
-) -> None:
-    recorder: Final = RecordingLogger()
-    result: Final = await execute(route, asynchronous, recording_server, {"callbacks": [recorder], "temperature": 0.25})
-    assert len(recording_server.requests) == 1
-    sent: Final = recording_server.requests[0]
-    body: Final = _OBJECT.validate_python(sent.body)
-    assert body["temperature"] == 0.25
-    assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body
-    if route == "chat":
-        assert isinstance(result, ModelResponse)
-        assert result.choices[0].message.content == "Hello from native Messages"
-        assert sent.path == "/v1/messages"
-    else:
-        assert isinstance(result, ResponsesAPIResponse)
-        assert result.output_text == "native response"
-        assert sent.path == "/responses"
-    success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event")
-    assert len(success) == 1
-    if isinstance(result, ModelResponse):
-        assert success[0].response is result
-    else:
-        logged: Final = success[0].response
-        assert isinstance(logged, ResponsesAPIResponse)
-        assert isinstance(result, ResponsesAPIResponse)
-        assert logged.id == result.id
-        assert logged.output_text == result.output_text
-
-
-@pytest.mark.asyncio
-async def test_native_inference_pre_call_edits_reach_the_provider(
-    route: Route, recording_server: RecordingServer
-) -> None:
-    class Edit(CustomLogger):
-        def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
-            request_body(kwargs)["temperature"] = 0.75
-
-    await execute(route, True, recording_server, {"callbacks": [Edit()]})
-    assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.75
-
-
-@pytest.mark.asyncio
-async def test_native_inference_does_not_restore_keywords_deleted_by_a_pre_call_hook(
-    route: Route, recording_server: RecordingServer
-) -> None:
-    class DropKeywords(CustomLogger):
-        async def async_pre_call_deployment_hook(
-            self, kwargs: dict[str, object], call_type: CallTypes | None
-        ) -> dict[str, object]:
-            return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")}
-
-    litellm.callbacks.append(DropKeywords())
-    await execute(
-        route, True, recording_server, {"temperature": 0.25, "extra_headers": {"x-deleted": "remove this header"}}
-    )
-
-    assert len(recording_server.requests) == 1
-    sent: Final = recording_server.requests[0]
-    assert "temperature" not in _OBJECT.validate_python(sent.body)
-    assert sent.headers.get("x-deleted") is None
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("from_credentials", (False, True))
-async def test_native_resource_setup_uses_deployment_hook_arguments(
-    route: Route,
-    recording_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-    from_credentials: bool,
-) -> None:
-    invalid_settings: Final = {"ssl_verify": object()}
-    credential: Final = CredentialItem(
-        credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
-    )
-    monkeypatch.setattr(litellm, "credential_list", [credential])
-
-    class Prepare(CustomLogger):
-        async def async_pre_call_deployment_hook(
-            self, kwargs: dict[str, object], call_type: CallTypes | None
-        ) -> dict[str, object]:
-            return {
-                **kwargs,
-                **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
-            }
-
-    litellm.callbacks.append(Prepare())
-    recorder: Final = RecordingLogger()
-    recording_server.expected_requests = 0
-    with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
-        await execute(route, True, recording_server, {"callbacks": [recorder]})
-    failure: Final = await recorder.wait_for_async("async_log_failure_event")
-    assert len(failure) == 1
-    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
-    assert not recording_server.requests
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
-    route: Route,
-    asynchronous: bool,
-    recording_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    monkeypatch.setattr(litellm, "max_budget", 1.0)
-    monkeypatch.setattr(litellm, "_current_cost", 2.0)
-    monkeypatch.setattr(litellm, "ssl_verify", object())
-    recorder: Final = RecordingLogger()
-    recording_server.expected_requests = 0
-    with pytest.raises(litellm.BudgetExceededError) as caught:
-        await execute(route, asynchronous, recording_server, {"callbacks": [recorder]})
-    failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
-    assert len(failure) == 1
-    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
-    assert not recording_server.requests
-    assert not any("success" in name for name in recorder.names)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_inference_provider_failure_is_terminal_and_shared_with_callbacks(
-    route: Route,
-    asynchronous: bool,
-    recording_server: RecordingServer,
-) -> None:
-    recorder: Final = RecordingLogger()
-    recording_server.enqueue(
-        ResponseSpec(body={"error": {"message": "slow down", "type": "rate_limit_error"}}, status=429)
-    )
-    with pytest.raises(RateLimitError) as caught:
-        await execute(route, asynchronous, recording_server, {"callbacks": [recorder]})
-    assert getattr(caught.value, "status_code", None) == 429
-    assert len(recording_server.requests) == 1
-    failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
-    assert len(failure) == 1
-    assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
-    assert not any("success" in name for name in recorder.names)
-
-
-@pytest.mark.asyncio
-async def test_unstarted_native_inference_has_no_provider_or_callback_effects(
-    route: Route,
-    recording_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    recorder: Final = RecordingLogger()
-    recording_server.expected_requests = 0
-    monkeypatch.setattr(litellm, "ssl_verify", object())
-    monkeypatch.setattr(litellm, "max_budget", 1.0)
-    monkeypatch.setattr(litellm, "_current_cost", 2.0)
-    pending: Final = native_call(route, True, recording_server, {"callbacks": [recorder]})
-    assert asyncio.iscoroutine(pending)
-    pending.close()
-    assert not recording_server.requests
-    assert not recorder.events
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "options",
-    (
-        {"stream": True},
-        {"extra_body": {"provider_option": True}},
-        {"mock_response": "mock"},
-        {"num_retries": 1},
-        {"use_chat_completions_api": True},
-        {"model_list": []},
-    ),
-)
-async def test_native_responses_declines_unsupported_requests_before_callbacks(
-    recording_server: RecordingServer,
-    options: Mapping[str, object],
-) -> None:
-    recorder: Final = RecordingLogger()
-    recording_server.expected_requests = 0
-    with pytest.raises(_native.RustBridgeDeclined):
-        native_call("responses", True, recording_server, {**options, "callbacks": [recorder]})
-    assert not recording_server.requests
-    assert not recorder.events
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_chat_validation_failure_is_terminal(
-    asynchronous: bool, recording_server: RecordingServer
-) -> None:
-    recording_server.expected_requests = 0
-    with pytest.raises(Exception, match="chat completions requires at least one message") as failure:
-        await execute("chat", asynchronous, recording_server, {"messages": []})
-    assert not isinstance(failure.value, _native.RustBridgeDeclined)
-    assert not recording_server.requests
-
-
-@pytest.mark.asyncio
-async def test_native_projection_reads_positional_parameters(route: Route, recording_server: RecordingServer) -> None:
-    from litellm.chat_completions.dispatch import (
-        _DISPATCH as chat_dispatch,  # pyright: ignore[reportPrivateUsage]  # exercise the request passed to the native boundary
-    )
-    from litellm.responses.dispatch import (
-        _DISPATCH as responses_dispatch,  # pyright: ignore[reportPrivateUsage]  # exercise the request passed to the native boundary
-    )
-
-    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE)
-    kwargs: Final = {"api_key": "test-key", "api_base": recording_server.base_url}
-    if route == "chat":
-        args: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.25)
-        request: Final = chat_dispatch.request(args, kwargs)
-        assert request is not None
-        await asyncio.to_thread(_native.completion, request)
-        assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.25
-    else:
-        response_args: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16)
-        response_request: Final = responses_dispatch.request(response_args, kwargs)
-        assert response_request is not None
-        await asyncio.to_thread(_native.responses, response_request)
-        body: Final = _OBJECT.validate_python(recording_server.requests[0].body)
-        assert body["instructions"] == "Be brief"
-        assert body["max_output_tokens"] == 16
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
-async def test_native_connection_settings_reach_the_provider(
-    route: Route,
-    asynchronous: bool,
-    source: str,
-    recording_server: RecordingServer,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    key: Final = "selected-key"
-    explicit: Final = source in ("explicit", "base_url")
-    monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
-    monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None))
-    monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None))
-    monkeypatch.setattr(
-        litellm,
-        "api_base",
-        None if source == "environment" else ("http://127.0.0.1:1" if explicit else recording_server.base_url),
-    )
-    monkeypatch.setenv(
-        "OPENAI_API_KEY" if route == "responses" else "ANTHROPIC_API_KEY", key if source == "environment" else "unused"
-    )
-    for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"):
-        monkeypatch.setenv(name, recording_server.base_url if source == "environment" else "http://127.0.0.1:1")
-    result: Final = await execute(
-        route,
-        asynchronous,
-        recording_server,
-        {
-            "api_key": key if explicit else ("" if source == "empty" else None),
-            "api_base": recording_server.base_url if source == "explicit" else ("" if source == "empty" else None),
-            **({"base_url": recording_server.base_url} if source == "base_url" else {}),
-        },
-    )
-    assert isinstance(result, ModelResponse | ResponsesAPIResponse)
-    assert len(recording_server.requests) == 1
-    headers: Final = recording_server.requests[0].headers
-    assert headers["x-api-key" if route == "chat" else "authorization"] == (key if route == "chat" else f"Bearer {key}")
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-@pytest.mark.parametrize("encoded", (False, True))
-async def test_native_responses_decode_continuation_ids(
-    asynchronous: bool, encoded: bool, recording_server: RecordingServer
-) -> None:
-    original: Final = "resp_upstream"
-    previous: Final = (
-        ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original)
-        if encoded
-        else original
-    )
-    await execute("responses", asynchronous, recording_server, {"previous_response_id": previous})
-    assert _OBJECT.validate_python(recording_server.requests[0].body)["previous_response_id"] == original
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_chat_uses_bound_positional_parameters(
-    asynchronous: bool, recording_server: RecordingServer
-) -> None:
-    recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
-    arguments: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.35)
-    supplied: Final = {"base_url": recording_server.base_url, "api_key": "test-key", "max_tokens": 32}
-    call: Final = chat_dispatch._DISPATCH.request(arguments, supplied)  # pyright: ignore[reportPrivateUsage]  # exercise the native request produced by public binding
-    assert call is not None
-    result: Final = (
-        await _native.acompletion(call) if asynchronous else await asyncio.to_thread(_native.completion, call)
-    )
-    assert isinstance(result, ModelResponse)
-    assert result.choices[0].message.content == "Hello from native Messages"
-    assert len(recording_server.requests) == 1
-    body: Final = _OBJECT.validate_python(recording_server.requests[0].body)
-    assert body["temperature"] == arguments[3]
-    assert body["max_tokens"] == supplied["max_tokens"]

tests/unit/rust_bridge/messages/test_secrets.py

@@ -14,8 +14,8 @@ from litellm.rust_bridge.host import settings
 from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES
 from litellm.rust_bridge.public_call import NativeCall
 from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
 from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
 
 pytest.importorskip("litellm.rust_bridge._native")
 

tests/unit/rust_bridge/ocr/test_secrets.py

@@ -16,8 +16,8 @@ from litellm.rust_bridge.host import settings
 from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
 from litellm.rust_bridge.public_call import NativeCall
 from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
+from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
 from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
 
 native: Final = pytest.importorskip("litellm.rust_bridge._native")