From 473e1a442ec93d58c50e435a83deb988d79f5331 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 21 Oct 2026 08:36:14 +0000 Subject: [PATCH] fix(rust-bridge): keep proxy /v1/messages calls on Rust when a cache or shared session is configured Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/messages/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 3 +-- litellm/rust_bridge/detectors.py | 4 ++-- .../messages/test_dispatch.py | 22 +++++++++++++++++-- .../responses/test_dispatch.py | 1 - tests/unit/rust_bridge/test_catalog.py | 6 +++++ tests/unit/rust_bridge/test_detectors.py | 2 +- 7 files changed, 30 insertions(+), 9 deletions(-) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cdb9a925403..59e02474c21 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -104,7 +104,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, - Feature.RESPONSE_CACHE: detectors.response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ebb221e98be..118dbf95818 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -169,7 +169,6 @@ MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, Feature.MESSAGES_INTERCEPTOR, - Feature.RESPONSE_CACHE, Feature.MODIFY_PARAMS, Feature.WORKLOAD_IDENTITY, Feature.CACHE_CONTROL_INJECTION, @@ -208,7 +207,7 @@ RULES: Final[Rules] = ( RouteRule(Route.OCR, Rollout.RUST_REQUIRED), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES), - RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "deepseek", "vertex_ai"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY, features=RESPONSES_PYTHON_FEATURES), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index 0ad1b787e38..2ed68f75c70 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -36,8 +36,8 @@ def mock_response(request: NativeCall) -> bool: def caller_client(request: NativeCall) -> bool: - """Python sends through the caller's HTTP client or shared session.""" - return present(request, "client") or present(request, "shared_session") + """Python sends through the caller's HTTP client. The proxy's shared session is only a pool native replaces.""" + return present(request, "client") def extra_body(request: NativeCall) -> bool: diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 50f71569d73..dbd4e5a0b3a 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -16,7 +16,7 @@ from litellm.messages.dispatch import ( ) from litellm.proxy import proxy_server from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target +from litellm.rust_bridge.catalog import MESSAGES_PYTHON_FEATURES, Feature, Route, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages from litellm.rust_bridge.public_call import NativeCall @@ -228,6 +228,25 @@ def test_only_requests_that_pass_through_to_a_messages_api_reach_a_messages_targ assert result is (native_result if passes_through else python_result) +def test_proxy_cache_and_shared_session_keep_messages_native(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "anthropic/claude-sonnet-4-5" + python_result: Final = response(model) + native_result: Final = response(model) + monkeypatch.setattr(litellm, "cache", Cache()) + result: Final = _DISPATCH.run( + (16, MESSAGES, model), + {"shared_session": object()}, + python=lambda *call_args, **call_kwargs: python_result, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ), + ) + assert result is native_result + + DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}] CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"} @@ -457,7 +476,6 @@ ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", " (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), (Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), - (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), ( diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 4498dbdf39d..19255379dcf 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -260,7 +260,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), (Feature.PROXY_MCP_TOOLS, {"tools": [{"type": "mcp", "server_label": "g", "server_url": "litellm_proxy"}]}, {}), (Feature.CALLER_CLIENT, {"client": object()}, {}), - pytest.param(Feature.CALLER_CLIENT, {"shared_session": object()}, {}, id="shared-session"), pytest.param(Feature.EXTRA_BODY, {"extra_body": {"stream": True}}, {}, id="extra-body-stream"), pytest.param( Feature.PROVIDER_AFFINITY, diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 704a4f73fe8..32690143fc4 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -128,3 +128,9 @@ def test_empty_constraints_match_nothing() -> None: assert catalog.decision( RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules ) is Decision.PYTHON + + +@pytest.mark.parametrize("provider", ("anthropic", "deepseek", "vertex_ai")) +def test_messages_passthrough_providers_opt_in_to_rust(provider: str) -> None: + context: Final = RouteContext(Route.MESSAGES, provider=provider, model="claude-sonnet-4-6", target=Target.MESSAGES) + assert catalog.rollout(context) is Rollout.RUST_OPT_IN diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 726128386d5..6c1553798c5 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -92,7 +92,7 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.prompt_management, {"model": "litellm_agent/gpt-5"}, True), (detectors.prompt_management, {"model": "gpt-5"}, False), (detectors.caller_client, {"client": object()}, True), - (detectors.caller_client, {"shared_session": object()}, True), + (detectors.caller_client, {"shared_session": object()}, False), (detectors.caller_client, {"client": None, "shared_session": None}, False), (detectors.caller_client, {}, False), (detectors.extra_body, {"extra_body": {"stream": True}}, True),