From 779441b47ba564696c35e86a66527e683c1fad31 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 00:30:19 +0000 Subject: [PATCH 01/35] fix(bedrock_mantle): source per-request AWS credential params from litellm_params when signing chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 26 ++++++- .../test_bedrock_mantle_transformation.py | 77 +++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 19 +++++ 3 files changed, 120 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..dfc234d08e0 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5,7 +5,7 @@ import ssl from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from contextlib import asynccontextmanager from functools import lru_cache -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints from urllib.parse import parse_qs, urlencode, urlparse, urlunparse @@ -20,6 +20,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -252,6 +253,24 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _aws_signing_overrides( + optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any] +) -> Mapping[str, Any]: + """AWS credential params for SigV4 signers that read them off optional_params. + + Only `bedrock`/`sagemaker` keep `aws_*` in optional_params: every other provider + spreads optional_params into the request body, so the params are stripped there + and survive on litellm_params alone. + """ + return MappingProxyType( + { + key: litellm_params[key] + for key in AWS_CREDENTIAL_KWARGS_KEYS + if optional_params.get(key) is None and litellm_params.get(key) is not None + } + ) + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -495,7 +514,10 @@ class BaseLLMHTTPHandler: headers, signed_json_body = provider_config.sign_request( headers=headers, - optional_params=optional_params, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, request_data=data, api_base=api_base, api_key=api_key, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 275fb460b9f..468cc9b9130 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -489,6 +489,83 @@ class TestBedrockMantleChatAuth: assert "/us-east-2/bedrock/aws4_request" in authorization assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws") + def test_completion_per_request_role_reaches_signer_and_not_the_body( + self, monkeypatch + ): + # Per-request aws_role_name/aws_session_name are stripped from optional_params + # for non-bedrock providers, so they must be sourced from litellm_params at + # signing time, and must never be serialized into the provider request body. + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + for var in ( + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + ): + monkeypatch.delenv(var, raising=False) + + credential_calls = [] + + def fake_get_credentials(self, **kwargs): + credential_calls.append(kwargs) + return Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + + monkeypatch.setattr(BaseAWSLLM, "get_credentials", fake_get_credentials) + + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data + requests.append({"headers": headers or {}, "body": json.loads(raw_body or "{}")}) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "google.gemma-4-31b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post + ): + litellm.completion( + model="bedrock_mantle/google.gemma-4-31b", + messages=[{"role": "user", "content": "hello"}], + aws_role_name="arn:aws:iam::000000000000:role/attributed-role", + aws_session_name="user-123", + aws_region_name="us-east-1", + ) + + assert len(credential_calls) == 1 + assert ( + credential_calls[0]["aws_role_name"] + == "arn:aws:iam::000000000000:role/attributed-role" + ) + assert credential_calls[0]["aws_session_name"] == "user-123" + assert requests[0]["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert not [key for key in requests[0]["body"] if key.startswith("aws_")] + class TestBedrockMantleProjectHeader: def test_validate_environment_sets_openai_project_header(self): diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..0906b39c514 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2071,3 +2071,22 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques retry_authorization = posts[1]["headers"]["Authorization"] assert retry_authorization.startswith("AWS4-HMAC-SHA256") assert retry_authorization != first_attempt_headers["Authorization"] + + +def test_aws_signing_overrides_only_fills_missing_credentials(): + from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides + + overrides = _aws_signing_overrides( + {"temperature": 0.2, "aws_region_name": "us-west-2"}, + { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + "api_key": "not-an-aws-param", + }, + ) + + assert dict(overrides) == { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + } From 2c7de60692d7a4fcd53964872d4355042d52b90b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 00:53:40 +0000 Subject: [PATCH 02/35] style: ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dfc234d08e0..cadf4c701e2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -253,9 +253,7 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False -def _aws_signing_overrides( - optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any] -) -> Mapping[str, Any]: +def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]: """AWS credential params for SigV4 signers that read them off optional_params. Only `bedrock`/`sagemaker` keep `aws_*` in optional_params: every other provider From 3275459aec0935b03b60950ba5d625854e2d6c9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:47:26 -0700 Subject: [PATCH 03/35] fix(mcp): cap tools preview and test-connection at the listing timeout and name the unreachable upstream --- .../mcp_server/rest_endpoints.py | 34 +++++--- .../mcp_server/test_rest_endpoints.py | 77 +++++++++++++++++-- 2 files changed, 96 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..d73277f4417 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal +import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -68,7 +70,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( ) -def _connection_error_message(exc: BaseException) -> str: +def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, TimeoutError): + return ( + f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " + "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." + ) if isinstance(exc, httpx.LocalProtocolError): return ( "Failed to connect to MCP server: a request header is malformed. " @@ -1136,6 +1144,7 @@ if MCP_AVAILABLE: mcp_auth_header: str | dict[str, str] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT, ) -> Mapping[str, object]: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1151,6 +1160,10 @@ if MCP_AVAILABLE: oauth2_headers: Headers extracted from the incoming request (may contain the litellm API key — must NOT be forwarded for M2M servers). raw_headers: Raw request headers forwarded for stdio env construction. + timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation* + combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB + timeouts) so an unreachable upstream yields this endpoint's JSON error + instead of an opaque load-balancer 504 with an empty body. Returns: The dict returned by *operation*, or an error dict on failure. @@ -1240,15 +1253,16 @@ if MCP_AVAILABLE: static_headers=request.static_headers, ) - client: Final = await global_mcp_server_manager._create_mcp_client( - server=server_model, - mcp_auth_header=mcp_auth_header, - extra_headers=merged_headers, - stdio_env=stdio_env, - cred_provider=preview_cred_provider, - ) + with anyio.fail_after(timeout_seconds): + client: Final = await global_mcp_server_manager._create_mcp_client( + server=server_model, + mcp_auth_header=mcp_auth_header, + extra_headers=merged_headers, + stdio_env=stdio_env, + cred_provider=preview_cred_provider, + ) - return await operation(client) + return await operation(client) except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise @@ -1257,7 +1271,7 @@ if MCP_AVAILABLE: return { "status": "error", "error": True, - "message": _connection_error_message(e), + "message": _connection_error_message(e, request.url, timeout_seconds), } async def _preview_openapi_tools(spec_path: str) -> dict: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..51b946c11b7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import sys from datetime import datetime @@ -13,6 +14,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, @@ -109,6 +111,71 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert "stack_trace" not in result + @pytest.mark.asyncio + async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def hanging_operation(client): + await asyncio.Event().wait() + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + @pytest.mark.asyncio + async def test_timeout_covers_client_creation(self, monkeypatch): + async def hanging_create_client(*args, **kwargs): + await asyncio.Event().wait() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + hanging_create_client, + ) + + async def unreached_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + def test_timeout_defaults_to_tool_listing_timeout(self): + default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default + assert default == MCP_TOOL_LISTING_TIMEOUT + + def test_connection_error_message_timeout_names_url_and_budget(self): + message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) + assert "https://api.example.com/mcp/" in message + assert "30s" in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. @@ -2881,17 +2948,17 @@ class TestConnectionErrorMessage: secret = "Bearer sk-super-secret-token" exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'") - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "header" in message.lower() assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed")) + message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out")) + message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): @@ -2901,11 +2968,11 @@ class TestConnectionErrorMessage: request=httpx.Request("POST", "http://x/"), response=response, ) - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message def test_unknown_error_falls_back_to_generic(self): - message = rest_endpoints._connection_error_message(RuntimeError("weird")) + message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message assert "proxy logs" in message.lower() From f4b5449c6a65cf167658fd5bb32695abda9633a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:48:04 -0700 Subject: [PATCH 04/35] fix(openai_like): strip cache_control ttl before forwarding /v1/messages to non-Anthropic providers --- litellm/llms/anthropic/common_utils.py | 33 +++++ litellm/llms/openai_like/README.md | 5 +- .../openai_like/messages/transformation.py | 38 +++++- ..._like_anthropic_messages_transformation.py | 121 ++++++++++++++++++ 4 files changed, 195 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 681a8397f66..695e3a313ef 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,6 +1302,39 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format +def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers + cache_type: Final = cache_control.get("type") + return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format + + +def _normalize_cache_control_value(value: object) -> object: + if isinstance(value, dict): + return normalize_cache_control_in_anthropic_payload(value) + if isinstance(value, list): + return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format + return value + + +def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers + """ + Return a copy of an Anthropic /v1/messages payload with every + ``cache_control`` entry reduced to ``{"type": }``, + recursing through message content blocks, system blocks, and tools. + + Anthropic itself accepts prompt-caching extensions such as ``ttl``, but + strict non-Anthropic implementations of the Messages API validate the field + literally and reject the whole request (``cache_control.ttl: 1h is not + supported``, ``cache_control.type is required``), which 400s clients like + Claude Code that always send cache hints. Non-dict ``cache_control`` values + are dropped entirely. The caller's payload is never mutated. + """ + return { # mutable-ok: JSON wire format, as sibling sanitizers + key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) + for key, value in payload.items() + if key != "cache_control" or isinstance(value, dict) + } + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e9aaafe48a1..e1409b81c35 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available. "constraints": { "temperature_max": 1.0, "temperature_min": 0.0, - "temperature_min_with_n_gt_1": 0.3 + "temperature_min_with_n_gt_1": 0.3, + // /v1/messages providers only: keep Anthropic cache_control extensions + // such as ttl instead of stripping them down to {"type": ...} + "cache_control_ttl": true }, // Optional: Special handling flags diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 11dc236064d..29973fe2101 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,11 +1,13 @@ from typing import Any, Final import litellm +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -19,7 +21,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control, thinking, tools, ...) is forwarded essentially unchanged to ``{api_base}/v1/messages``, so Anthropic-only features that the - Anthropic->OpenAI translation would otherwise drop are preserved. Response + Anthropic->OpenAI translation would otherwise drop are preserved. The one + exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) + are stripped unless ``supports_cache_control_ttl`` says otherwise. Response parsing and streaming are inherited from the native Anthropic config. """ @@ -53,6 +57,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def supports_cache_control_ttl(self) -> bool: + return False + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict, # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: matches dict-typed base signature + ) -> dict: # mutable-ok: matches dict-typed base signature + """ + Anthropic ignores prompt-caching hints it cannot honor, but strict + non-Anthropic implementations of the Messages API 400 the whole request + on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h + is not supported``), so unless the provider declares ttl support the + hints are reduced to their portable ``{"type": ...}`` core. + """ + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if self.supports_cache_control_ttl(): + return request + return normalize_cache_control_in_anthropic_payload(request) + def get_complete_url( self, api_base: str | None, @@ -91,6 +124,9 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True + def supports_cache_control_ttl(self) -> bool: + return bool(self._provider.constraints.get("cache_control_ttl")) + def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 33e677b000e..2cdd969b00f 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -317,3 +317,124 @@ def test_json_provider_messages_config_probes_capabilities_under_provider_slug() ) assert JSONProviderAnthropicMessagesConfig(provider).custom_llm_provider == "exampleprovider" assert OpenAILikeAnthropicMessagesConfig().custom_llm_provider == "anthropic" + + +def _cache_control_request_params() -> tuple[list, dict]: + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "write a regex for a US phone number", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + optional_params = { + "max_tokens": 256, + "system": [ + { + "type": "text", + "text": "You are Claude Code.", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + return messages, optional_params + + +def test_request_strips_cache_control_ttl_everywhere(config): + """Regression: Claude Code always sends ``cache_control: {type: ephemeral, + ttl: 1h}``, and strict non-Anthropic /v1/messages validators 400 the whole + request on the ttl extension (``cache_control.ttl: 1h is not supported``).""" + messages, optional_params = _cache_control_request_params() + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_defaults_missing_cache_control_type_and_drops_non_dict(config): + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "a", "cache_control": {"ttl": "1h"}}, + {"type": "text", "text": "b", "cache_control": None}, + ], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + blocks = payload["messages"][0]["content"] + assert blocks[0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in blocks[1] + + +def test_native_anthropic_config_keeps_cache_control_ttl(): + """Anthropic itself accepts ttl, so the normalization must stay scoped to + the OpenAI-like passthrough and never reach the native Anthropic path.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + messages, optional_params = _cache_control_request_params() + payload = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-20250514", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +def test_json_provider_constraint_opts_into_cache_control_ttl(): + from litellm.llms.openai_like.json_loader import SimpleProviderConfig + from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, + ) + + base_data = {"base_url": "https://api.example.com/v1", "api_key_env": "EXAMPLE_API_KEY"} + strict = JSONProviderAnthropicMessagesConfig(SimpleProviderConfig(slug="strictprov", data=base_data)) + lenient = JSONProviderAnthropicMessagesConfig( + SimpleProviderConfig(slug="lenientprov", data={**base_data, "constraints": {"cache_control_ttl": True}}) + ) + + def transform(provider_config): + messages, optional_params = _cache_control_request_params() + return provider_config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} From d3268e4e184f8b7cde6cce4473e14b1acedfc751 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:26:04 -0700 Subject: [PATCH 05/35] test(mcp): wrap over-long connection error message calls --- .../proxy/_experimental/mcp_server/test_rest_endpoints.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 51b946c11b7..e66101f6177 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -2954,11 +2954,15 @@ class TestConnectionErrorMessage: assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectTimeout("timed out"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): From 0baf376efd691e139902ca7933b83666ea459a41 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:07:05 -0700 Subject: [PATCH 06/35] fix(openai_like): scope cache_control normalization to Messages API locations Rewrite the sanitizer without recursion (the code-quality gate rejects new recursive functions) and only touch cache_control where the Messages API defines it: the request, system blocks, tools, message content blocks, and tool_result content. Application data such as tool_use.input and tool input_schema is left untouched even when it contains a cache_control key --- litellm/llms/anthropic/common_utils.py | 78 +++++++++++++++---- ..._like_anthropic_messages_transformation.py | 62 +++++++++++++++ 2 files changed, 124 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 695e3a313ef..1ef14362601 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,37 +1302,83 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format -def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers +def _normalized_cache_control(cache_control: object) -> dict[str, str] | None: # mutable-ok: JSON wire format + if not isinstance(cache_control, Mapping): + return None cache_type: Final = cache_control.get("type") return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format -def _normalize_cache_control_value(value: object) -> object: - if isinstance(value, dict): - return normalize_cache_control_in_anthropic_payload(value) - if isinstance(value, list): - return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format - return value +def _with_portable_cache_control(block: Mapping[str, object]) -> dict[str, object]: # mutable-ok: JSON wire format + if "cache_control" not in block: + return dict(block) # mutable-ok: JSON wire format + normalized: Final = _normalized_cache_control(block["cache_control"]) + rest: Final = {key: value for key, value in block.items() if key != "cache_control"} # mutable-ok: JSON wire format + return rest if normalized is None else {**rest, "cache_control": normalized} # mutable-ok: JSON wire format -def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers +def _with_portable_cache_control_in_blocks(blocks: object) -> object: + if isinstance(blocks, str) or not isinstance(blocks, Sequence): + return blocks + return [ # mutable-ok: JSON wire format + _with_portable_cache_control(block) if isinstance(block, Mapping) else block for block in blocks + ] + + +def _with_portable_cache_control_in_content_block(block: object) -> object: + if not isinstance(block, Mapping): + return block + portable: Final = _with_portable_cache_control(block) + if portable.get("type") != "tool_result" or "content" not in portable: + return portable + return { # mutable-ok: JSON wire format + **portable, + "content": _with_portable_cache_control_in_blocks(portable["content"]), + } + + +def _with_portable_cache_control_in_message(message: object) -> object: + if not isinstance(message, Mapping) or "content" not in message: + return message + content: Final = message["content"] + if isinstance(content, str) or not isinstance(content, Sequence): + return message + return { # mutable-ok: JSON wire format + **message, + "content": [_with_portable_cache_control_in_content_block(block) for block in content], + } + + +def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format + payload: Mapping[str, object], +) -> dict[str, object]: """ Return a copy of an Anthropic /v1/messages payload with every - ``cache_control`` entry reduced to ``{"type": }``, - recursing through message content blocks, system blocks, and tools. + ``cache_control`` entry reduced to ``{"type": }`` + at the places the Messages API defines it: the request itself, system + blocks, tools, message content blocks, and ``tool_result`` content blocks. + Application data such as ``tool_use.input`` and tool ``input_schema`` is + never touched, even when it happens to contain a ``cache_control`` key. Anthropic itself accepts prompt-caching extensions such as ``ttl``, but strict non-Anthropic implementations of the Messages API validate the field literally and reject the whole request (``cache_control.ttl: 1h is not supported``, ``cache_control.type is required``), which 400s clients like - Claude Code that always send cache hints. Non-dict ``cache_control`` values - are dropped entirely. The caller's payload is never mutated. + Claude Code that send cache hints. Non-dict ``cache_control`` values are + dropped entirely. The caller's payload is never mutated. """ - return { # mutable-ok: JSON wire format, as sibling sanitizers - key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) - for key, value in payload.items() - if key != "cache_control" or isinstance(value, dict) + portable: Final = _with_portable_cache_control(payload) + scoped: Final = { # mutable-ok: JSON wire format + key: ( + _with_portable_cache_control_in_blocks(value) + if key in ("system", "tools") + else [_with_portable_cache_control_in_message(message) for message in value] + if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) + else value + ) + for key, value in portable.items() } + return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 2cdd969b00f..d325492914e 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -438,3 +438,65 @@ def test_json_provider_constraint_opts_into_cache_control_ttl(): assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(config): + """Regression: the sanitizer must only touch ``cache_control`` where the + Messages API defines it (request, system, tools, content blocks, tool_result + content), never application data such as ``tool_use.input`` or a tool's + ``input_schema`` that happens to contain a ``cache_control`` key.""" + tool_input = {"cache_control": {"type": "ephemeral", "ttl": "1h"}, "query": "x"} + input_schema = { + "type": "object", + "properties": {"cache_control": {"type": "string", "ttl": "1h"}}, + } + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": tool_input}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "content": [ + {"type": "text", "text": "result", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"type": "text", "text": "plain string content stays", "cache_control": {"ttl": "1h"}}, + ], + }, + {"role": "user", "content": "a plain string message"}, + ] + optional_params = { + "max_tokens": 64, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "tools": [ + { + "name": "lookup", + "input_schema": input_schema, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["input_schema"] == input_schema + assert payload["messages"][0]["content"][0]["input"] == tool_input + tool_result = payload["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"} + assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} From c32eb41aad3b7b087c7d0023a71876d3dea6511d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:17:13 -0700 Subject: [PATCH 07/35] feat(openai_like): let a passthrough deployment keep cache_control ttl via model_info.cache_control_ttl The supported_endpoints passthrough had no way to keep ttl for an upstream that honors it, so the deployment now opts in with model_info.cache_control_ttl: true, injected into the config the same way the providers.json constraint is for JSON providers --- .../messages/handler.py | 8 +- .../openai_like/messages/transformation.py | 16 +-- ...erimental_pass_through_messages_handler.py | 97 ++++++++++--------- ..._like_anthropic_messages_transformation.py | 17 ++++ 4 files changed, 84 insertions(+), 54 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 69985bcdaa3..b82903d6f87 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -99,6 +99,10 @@ def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: return isinstance(supported_endpoints, (list, tuple)) and "/v1/messages" in supported_endpoints +def _deployment_supports_cache_control_ttl(model_info: object) -> bool: + return isinstance(model_info, dict) and model_info.get("cache_control_ttl") is True + + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -568,7 +572,9 @@ def anthropic_messages_handler( OpenAILikeAnthropicMessagesConfig, ) - anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig() + anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig( + cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")), + ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if _should_route_to_responses_api(custom_llm_provider, original_model, model): diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 29973fe2101..ac99617521c 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -23,10 +23,15 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``{api_base}/v1/messages``, so Anthropic-only features that the Anthropic->OpenAI translation would otherwise drop are preserved. The one exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) - are stripped unless ``supports_cache_control_ttl`` says otherwise. Response - parsing and streaming are inherited from the native Anthropic config. + are stripped unless the deployment opts in with + ``model_info.cache_control_ttl: true``. Response parsing and streaming are + inherited from the native Anthropic config. """ + def __init__(self, cache_control_ttl: bool = False) -> None: + super().__init__() + self._cache_control_ttl: Final = cache_control_ttl + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -58,7 +63,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): return False def supports_cache_control_ttl(self) -> bool: - return False + return self._cache_control_ttl def transform_anthropic_messages_request( self, @@ -114,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): """ def __init__(self, provider: SimpleProviderConfig): - super().__init__() + super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl"))) self._provider = provider @property @@ -124,9 +129,6 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True - def supports_cache_control_ttl(self) -> bool: - return bool(self._provider.constraints.get("cache_control_ttl")) - def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ad4c3d6bfbb..e819433c269 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -296,21 +296,15 @@ async def test_bedrock_converse_budget_tokens_preserved(): mock_acompletion.assert_called_once() call_kwargs = mock_acompletion.call_args.kwargs - print( - "acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str) - ) + print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) # Verify thinking parameter is passed through with budget_tokens preserved thinking_param = call_kwargs.get("thinking") - assert ( - thinking_param is not None - ), "thinking parameter should be passed to acompletion" - assert ( - thinking_param.get("type") == "enabled" - ), "thinking.type should be 'enabled'" - assert ( - thinking_param.get("budget_tokens") == 1024 - ), f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + assert thinking_param is not None, "thinking parameter should be passed to acompletion" + assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_param.get("budget_tokens") == 1024, ( + f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + ) def test_openai_model_with_thinking_converts_to_reasoning(): @@ -342,23 +336,18 @@ def test_openai_model_with_thinking_converts_to_reasoning(): call_kwargs = mock_responses.call_args.kwargs # Verify reasoning is set (converted from thinking) - assert ( - "reasoning" in call_kwargs - ), "reasoning should be passed to litellm.responses" + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" # budget_tokens=1024 -> effort="low" (at the LOW budget threshold) # reasoning_auto_summary is False by default, so no summary key expected_reasoning = {"effort": "low"} assert call_kwargs["reasoning"] == expected_reasoning, ( - f"reasoning should be {expected_reasoning} for budget_tokens=1024, " - f"got {call_kwargs.get('reasoning')}" + f"reasoning should be {expected_reasoning} for budget_tokens=1024, got {call_kwargs.get('reasoning')}" ) assert "summary" not in call_kwargs["reasoning"] # Verify thinking is NOT passed directly to the Responses API - assert ( - "thinking" not in call_kwargs - ), "thinking should NOT be passed directly to litellm.responses" + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: @@ -411,9 +400,7 @@ class TestThinkingParameterTransformation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "detailed"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "detailed"}} finally: litellm.reasoning_auto_summary = original @@ -611,9 +598,9 @@ class TestThinkingSummaryPreservation: mock_responses.assert_called_once() call_kwargs = mock_responses.call_args.kwargs reasoning = call_kwargs["reasoning"] - assert ( - reasoning["summary"] == "concise" - ), f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + assert reasoning["summary"] == "concise", ( + f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + ) def test_responses_adapter_preserves_summary(self): """translate_thinking_to_reasoning should include summary when user provides it.""" @@ -622,9 +609,7 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -638,11 +623,7 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: @@ -659,9 +640,7 @@ class TestThinkingSummaryPreservation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "concise"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "concise"}} def test_translate_thinking_for_model_disabled_stays_plain_string_when_auto_summary_enabled(self): """Disabled thinking must stay a plain string even when reasoning_auto_summary is on.""" @@ -807,9 +786,7 @@ def test_presanitized_flag_not_leaked_to_provider_params(): def fake_base_handler(*args, **kwargs): captured.update(kwargs) - captured["optional"] = kwargs.get( - "anthropic_messages_optional_request_params", {} - ) + captured["optional"] = kwargs.get("anthropic_messages_optional_request_params", {}) return "stub" with patch.object( @@ -974,6 +951,38 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert "config" not in captured +@pytest.mark.parametrize( + "model_info, expected_ttl_support", + [ + ({"supported_endpoints": ["/v1/messages"]}, False), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, True), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": "yes"}, False), + ], +) +def test_gate_passthrough_forwards_cache_control_ttl_only_when_deployment_opts_in( + monkeypatch, model_info, expected_ttl_support +): + """The passthrough config strips cache_control.ttl unless the deployment sets + model_info.cache_control_ttl to exactly true.""" + from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, + ) + + captured, _ = _gate_stubs(monkeypatch) + + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="openai/some-model", + api_key="sk-test", + api_base="https://host/v1", + model_info=model_info, + ) + + assert result == "native-passthrough" + assert captured["config"].supports_cache_control_ttl() is expected_ttl_support + + def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): """Regional and provider-prefixed Claude 4.8+/5 entries carry ``supports_mid_conversation_system``, but the bare first-party keys @@ -987,9 +996,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys import litellm - cost_map_path = os.path.join( - os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" - ) + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(cost_map_path) as f: cost_map = json.load(f) rules = cost_map["fallback_generalizations"]["rules"] @@ -1028,9 +1035,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), ], ) -async def test_messages_strips_provider_prefix_exactly_once( - requested_model, expected_wire_model, expected_url -): +async def test_messages_strips_provider_prefix_exactly_once(requested_model, expected_wire_model, expected_url): """ BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index d325492914e..e33b03afdff 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -414,6 +414,23 @@ def test_native_anthropic_config_keeps_cache_control_ttl(): assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} +def test_deployment_opt_in_keeps_cache_control_ttl(): + config = OpenAILikeAnthropicMessagesConfig(cache_control_ttl=True) + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 16}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + def test_json_provider_constraint_opts_into_cache_control_ttl(): from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.llms.openai_like.messages.transformation import ( From 692f3b513ca9ac2fafb225c24d58f7ee5152eae6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:43:01 -0700 Subject: [PATCH 08/35] fix(proxy-extras): recover the v2 migration resolver from concurrent migrate deploy deadlocks Two instances racing prisma migrate deploy on one database deadlock on CREATE INDEX CONCURRENTLY: the victim gets P3018 with 40P01 and the survivor then sees the failed ledger row as P3009. Both were treated as unrecoverable, so neither instance came up. Roll the deadlocked migration's ledger row back and retry the deploy on P3018, consult the failed row's logs in _prisma_migrations to do the same on P3009, and retry a deadlock reported without a Prisma error code. Genuinely broken migrations still fail fast. --- .../litellm_proxy_extras/utils.py | 91 ++++++++++- .../tests/test_setup_database_fail_fast.py | 141 ++++++++++++++++++ 2 files changed, 228 insertions(+), 4 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..fb948afd200 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -40,6 +40,8 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") +_MIGRATION_DEADLOCK_MARKER = "deadlock detected" + _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile( r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE @@ -262,6 +264,48 @@ class ProxyExtrasDBManager: env=prisma_env, ) + @staticmethod + def _roll_back_migration_best_effort(migration_name: str) -> None: + """Mark a migration rolled back, tolerating a concurrent resolver + having already done it.""" + try: + ProxyExtrasDBManager._roll_back_migration(migration_name) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + + @staticmethod + def _failed_migration_logs(migration_name: str) -> str: + """Logs recorded on the migration's failed _prisma_migrations row. + + P3009 stderr does not carry the original failure, so this is the only + way to tell a migration that lost a deadlock race against a concurrent + migrate deploy from one whose SQL is genuinely broken. Returns "" when + psycopg is missing, the DB is unreachable, or no failed row exists. + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + return "" + + try: + import psycopg + except ImportError: + return "" + + cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + try: + with psycopg.connect( + cleaned_url, connect_timeout=10, autocommit=True + ) as conn: + row = conn.execute( + "SELECT logs FROM _prisma_migrations " + "WHERE migration_name = %s AND finished_at IS NULL " + "AND rolled_back_at IS NULL", + (migration_name,), + ).fetchone() + except (psycopg.OperationalError, psycopg.DatabaseError): + return "" + return (row[0] or "") if row else "" + @staticmethod def _resolve_specific_migration(migration_name: str): """Mark a specific migration as applied""" @@ -658,7 +702,8 @@ class ProxyExtrasDBManager: v2 migration resolver (opt-in via --use_v2_migration_resolver). Runs `prisma migrate deploy` and handles standard recovery paths - (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does + (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a + concurrent migrate deploy). Critically, it does NOT call `_resolve_all_migrations` — the diff-and-force recovery that caused schema thrashing when two LiteLLM versions contended for the same DB during rolling deploys. @@ -764,6 +809,22 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue + if migration_match and _MIGRATION_DEADLOCK_MARKER in ( + ProxyExtrasDBManager._failed_migration_logs( + migration_match.group(1) + ) + ): + logger.info( + "Migration %s lost a deadlock race against a " + "concurrent migrate deploy, rolling its ledger " + "row back and retrying", + migration_match.group(1), + ) + ProxyExtrasDBManager._roll_back_migration_best_effort( + migration_match.group(1) + ) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" @@ -809,11 +870,33 @@ class ProxyExtrasDBManager: ) from resolve_err continue + if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "Migration %s deadlocked against a concurrent " + "migrate deploy, rolling its ledger row back " + "and retrying", + migration_match.group(1), + ) + ProxyExtrasDBManager._roll_back_migration_best_effort( + migration_match.group(1) + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e + if _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "prisma migrate deploy attempt %s deadlocked against " + "a concurrent migrate deploy, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" @@ -821,9 +904,9 @@ class ProxyExtrasDBManager: raise RuntimeError( "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts or repeated idempotent-recovery " - "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + "exhausted by timeouts, deadlock retries, or repeated " + "idempotent-recovery continues). Check database connectivity, " + "load, and _prisma_migrations ledger state." ) finally: os.chdir(original_dir) diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 8d66bf872de..c4347a91dce 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -240,3 +240,144 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" + + +_DEADLOCK_P3018_STDERR = ( + "Error: P3018\n" + "Migration name: 20260415120000_health_check_latest_per_model_index\n" + "Database error code: 40P01\n" + "deadlock detected" +) + + +def _stub_v2_env(monkeypatch, tmp_path): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda _: None) + + +def _succeed_after(failures: int, stderr: str): + calls = {"n": 0} + + class _OkResult: + stdout = "Applied migration.\n" + stderr = "" + + def _run(*args, **kwargs): + if "deploy" not in args[0]: + return _OkResult() + calls["n"] += 1 + if calls["n"] <= failures: + raise subprocess.CalledProcessError( + returncode=1, cmd=args[0], stderr=stderr, output="" + ) + return _OkResult() + + return _run + + +def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: losing the migrate deploy deadlock race against a concurrent + instance rolls the ledger row back and retries instead of dying.""" + _stub_v2_env(monkeypatch, tmp_path) + + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): + """v2: a deadlock on every attempt still fails after the retry budget.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) + + with patch( + "subprocess.run", + side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR), + ): + with pytest.raises(RuntimeError, match="after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: the surviving instance sees the victim's failed ledger row as P3009. + When that row's logs show a deadlock, roll it back and retry.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_failed_migration_logs", + lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock", + ) + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): + """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260101000000_genuinely_broken` migration started at " + "2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_failed_migration_logs", + lambda name: 'ERROR: syntax error at or near "BRKN"', + ) + + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): + """v2: a deadlock reported without a Prisma error code (the advisory-lock + waiter as victim) is retried, not fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr( + "subprocess.run", _succeed_after(1, "Database error: deadlock detected") + ) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True From 5192b2162c987f260c9c33700343a73ea4676749 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:04:12 -0700 Subject: [PATCH 09/35] fix(proxy-extras): schema-qualify the _prisma_migrations logs lookup for non-public Prisma schemas --- litellm-proxy-extras/litellm_proxy_extras/utils.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index fb948afd200..97b6b1c667c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -292,14 +292,22 @@ class ProxyExtrasDBManager: return "" cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + ledger_table = psycopg.sql.SQL("{}.{}").format( + psycopg.sql.Identifier( + ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" + ), + psycopg.sql.Identifier("_prisma_migrations"), + ) try: with psycopg.connect( cleaned_url, connect_timeout=10, autocommit=True ) as conn: row = conn.execute( - "SELECT logs FROM _prisma_migrations " - "WHERE migration_name = %s AND finished_at IS NULL " - "AND rolled_back_at IS NULL", + psycopg.sql.SQL( + "SELECT logs FROM {} " + "WHERE migration_name = %s AND finished_at IS NULL " + "AND rolled_back_at IS NULL" + ).format(ledger_table), (migration_name,), ).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): From 1cd99a036e0538bb61b280e17639545c75374d81 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:29:58 -0700 Subject: [PATCH 10/35] fix(router): route Claude Code subagents through session router --- litellm/router.py | 80 +++++++++++++++++- tests/test_litellm/test_router.py | 129 ++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..b49269f2457 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -353,6 +353,8 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: @@ -12546,6 +12548,77 @@ class Router: return None return candidates[0] + @staticmethod + def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: + proxy_server_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + headers: Final = proxy_server_request.get("headers") + if not isinstance(headers, Mapping): + return None + return next( + ( + value + for key, value in headers.items() + if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) + ), + None, + ) + + def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: + session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") + if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: + return None + metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata: Final = request_kwargs.get(metadata_name) + if not isinstance(metadata, Mapping): + return None + caller_scope: Final = metadata.get("user_api_key_hash") + if not isinstance(caller_scope, str) or not caller_scope: + return None + return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + + async def _resolve_claude_code_session_router( + self, + model: str, + registered_model_name: str, + request_kwargs: Mapping[str, object], + ) -> str: + cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) + if cache_key is None or not isinstance(request_kwargs, dict): + return registered_model_name + + agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self.cache.async_get_cache(key=cache_key) + if not isinstance(bound_model, str): + return registered_model_name + bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model + if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: + await self.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.cache.async_set_cache( + key=cache_key, + value=bound_model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) + return bound_registered_model + + if self._request_header(request_kwargs, "x-app") != "cli": + return registered_model_name + if request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name + if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: + await self.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.cache.async_set_cache( + key=cache_key, + value=model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + return registered_model_name + async def async_pre_routing_hook( self, model: str, @@ -12565,7 +12638,12 @@ class Router: the alias, since spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ - registered_model_name: Final = self._get_model_from_alias(model=model) or model + requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model + registered_model_name: Final = await self._resolve_claude_code_session_router( + model=model, + registered_model_name=requested_registered_model_name, + request_kwargs=request_kwargs, + ) ######################################################### # Run the routing-plugin pipeline, if any plugins are configured. diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..d22a1cc04ca 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8316,6 +8316,135 @@ class TestConsumedRequestTagsStamp: assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"] +class TestClaudeCodeSubagentSessionRouterBinding: + class _RewriteStrategy: + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse( + model="cheap-model", + messages=messages, + routing_decision={ + "router_model_name": "smart-router", + "router_type": "complexity", + "routed_model": "cheap-model", + "cause": "heuristic_scorer", + }, + ) + + @classmethod + def _router(cls) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + router = litellm.Router( + model_list=[ + { + "model_name": "cheap-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "cheap response"}, + }, + { + "model_name": "expensive-model", + "litellm_params": {"model": "openai/gpt-4o", "mock_response": "expensive response"}, + }, + ] + ) + router.complexity_routers = { + "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + } + return router + + @staticmethod + def _request_kwargs( + *, + key_hash: str = "key-hash-a", + app: str = "cli", + agent_id: str | None = None, + fallback_depth: int | None = None, + ) -> dict: + headers = { + "X-Claude-Code-Session-Id": "session-1234", + "x-app": app, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + } + return { + "metadata": {"user_api_key_hash": key_hash}, + "proxy_server_request": {"headers": headers}, + **({"fallback_depth": fallback_depth} if fallback_depth is not None else {}), + } + + @pytest.mark.asyncio + async def test_subagent_concrete_model_uses_the_main_sessions_router(self): + router = self._router() + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "cheap response" + assert subagent_kwargs["metadata"]["model_group"] == "smart-router" + assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" + + @pytest.mark.asyncio + async def test_main_direct_model_clears_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook(model="expensive-model", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is None + + @pytest.mark.asyncio + async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(app="cli-bg"), + ) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(fallback_depth=1), + ) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is not None + assert response.model == "cheap-model" + + @pytest.mark.asyncio + async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(key_hash="key-hash-b", agent_id="agent-1234"), + ) + + assert response is None + + class TestAutoRouterMaxInputCharsWiring: """`auto_router_max_input_chars` on the deployment has to reach the AutoRouter that embeds prompts. From 82d046c42821ca3b24036dd57e362252564c0596 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:43:48 -0700 Subject: [PATCH 11/35] fix(router): make Claude session cleanup best effort --- litellm/router.py | 14 ++++++++++++-- tests/test_litellm/test_router.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b49269f2457..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,6 +12578,16 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + try: + await self.cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; " + "the binding may remain until its TTL expires: %s", + e, + ) + async def _resolve_claude_code_session_router( self, model: str, @@ -12595,7 +12605,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12610,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d22a1cc04ca..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8409,6 +8409,25 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None + @pytest.mark.asyncio + async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(), + ) + + assert response is None + redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From e3a61c82da9f8dbe42fdfcc4907af7b3a2901392 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:46:41 -0700 Subject: [PATCH 12/35] test(router): register indirect session routing coverage --- tests/code_coverage_tests/router_code_coverage.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a5e00799519..60b56b7fac6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -82,6 +82,10 @@ ignored_function_names = [ "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call "_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name) + "_request_header", # Tested through Claude Code session routing in test_router.py + "_claude_code_session_router_cache_key", # Tested through Claude Code session routing in test_router.py + "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py + "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py ] From 2d4301589c1e741489f68e6eef3c2d112da91e2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:10:31 +0000 Subject: [PATCH 13/35] fix(router): keep serving when Claude Code session router cleanup fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 12 ++++-------- tests/test_litellm/test_router.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..d6d9f20085f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,15 +12578,11 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + async def _clear_claude_code_session_router(self, cache_key: str) -> None: try: await self.cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; " - "the binding may remain until its TTL expires: %s", - e, - ) + verbose_router_logger.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) async def _resolve_claude_code_session_router( self, @@ -12605,7 +12601,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12620,7 +12616,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..b66dbf6aa7d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,24 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_main_direct_model_still_served_when_cache_delete_fails(self): + router = self._router() + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() + ) + + async def failing_delete(key: str) -> None: + raise Exception("Redis circuit breaker is open — skipping async_delete_cache") + + router.cache.async_delete_cache = failing_delete + + response = await router.acompletion( + model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() + ) + + assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 46502f58042a41619648be7acc67108079f737e2 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:11:23 +0000 Subject: [PATCH 14/35] Revert "fix(router): keep serving when Claude Code session router cleanup fails" This reverts commit 2d4301589c1e741489f68e6eef3c2d112da91e2a. --- litellm/router.py | 12 ++++++++---- tests/test_litellm/test_router.py | 18 ------------------ 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d6d9f20085f..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,11 +12578,15 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - async def _clear_claude_code_session_router(self, cache_key: str) -> None: + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: await self.cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; " + "the binding may remain until its TTL expires: %s", + e, + ) async def _resolve_claude_code_session_router( self, @@ -12601,7 +12605,7 @@ class Router: return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12616,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b66dbf6aa7d..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,24 +8428,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() - @pytest.mark.asyncio - async def test_main_direct_model_still_served_when_cache_delete_fails(self): - router = self._router() - await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() - ) - - async def failing_delete(key: str) -> None: - raise Exception("Redis circuit breaker is open — skipping async_delete_cache") - - router.cache.async_delete_cache = failing_delete - - response = await router.acompletion( - model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() - ) - - assert response.choices[0].message.content == "expensive response" - @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 8a0967443d84ecc02c10caf6ca55385b907b11b2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 20:38:13 -0700 Subject: [PATCH 15/35] fix(router): isolate Claude session binding cache --- litellm/router.py | 19 +++++++++++++------ tests/test_litellm/test_router.py | 14 ++++++++++++++ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..fb8625bc6ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -782,6 +782,10 @@ class Router: self.cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._claude_code_session_router_cache: DualCache = DualCache( + redis_cache=redis_cache, + in_memory_cache=InMemoryCache(), + ) ### SCHEDULER ### self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) @@ -1102,8 +1106,8 @@ class Router: ``` and caching to just work. """ - if self.cache.redis_cache is None: - self.cache.redis_cache = cache + self.cache.attach_redis_cache(cache) + self._claude_code_session_router_cache.attach_redis_cache(cache) # Maps a routing strategy string to the attribute on `self` that holds # the default group's strategy selector for that strategy. (The selectors @@ -12580,7 +12584,7 @@ class Router: async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: - await self.cache.async_delete_cache(key=cache_key) + await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request verbose_router_logger.warning( "Failed to delete Claude Code session router binding; " @@ -12600,14 +12604,14 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self.cache.async_get_cache(key=cache_key) + bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=cache_key) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name - await self.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=bound_model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -12622,7 +12626,7 @@ class Router: if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name - await self.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -13450,6 +13454,9 @@ class Router: def flush_cache(self): litellm.cache = None self.cache.flush_cache() + session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache + if session_in_memory_cache is not None: + session_in_memory_cache.flush_cache() def reset(self): ## clean up on close diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..a528266d738 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,20 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_session_bindings_do_not_evict_router_rate_limit_state(self): + router = self._router() + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 1 + + for session_index in range(201): + request_kwargs = self._request_kwargs() + request_kwargs["proxy_server_request"]["headers"]["X-Claude-Code-Session-Id"] = ( + f"session-{session_index:04d}" + ) + await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs) + + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 2 + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 6adc14b4b12250ddc927682c8153469976e87888 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 21:13:45 -0700 Subject: [PATCH 16/35] fix(router): preserve Claude subagent fallbacks --- litellm/router.py | 4 ++-- tests/test_litellm/test_router.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fb8625bc6ba..6f1f1bc700b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12601,6 +12601,8 @@ class Router: cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) if cache_key is None or not isinstance(request_kwargs, dict): return registered_model_name + if request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: @@ -12621,8 +12623,6 @@ class Router: if self._request_header(request_kwargs, "x-app") != "cli": return registered_model_name - if request_kwargs.get("fallback_depth") not in (None, 0): - return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a528266d738..b413bb18f04 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8464,6 +8464,19 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is not None assert response.model == "cheap-model" + @pytest.mark.asyncio + async def test_subagent_fallback_does_not_reapply_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234", fallback_depth=1), + ) + + assert response is None + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From 7b919f89a85ea7671fd0c3ac2fb27d31fb74201b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 10:00:47 -0700 Subject: [PATCH 17/35] fix(router): track routed model in fallback attempts --- .../router_utils/fallback_event_handlers.py | 5 +-- .../test_fallback_event_handlers.py | 21 +++++++++++ tests/test_litellm/test_router.py | 35 +++++++++++++++++-- 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3d37ca216a7..0167721f9fe 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -470,10 +470,11 @@ async def run_async_fallback( attempted: Final = ( carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() ) - attempted.record(original_model_group) + failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group + attempted.record(failed_model_group) for mg in fallback_model_group: - if mg == original_model_group: + if mg == failed_model_group: continue if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: verbose_router_logger.info( diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 894b2d9e74f..9e51a60364b 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -614,6 +614,27 @@ async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call ) +@pytest.mark.asyncio +async def test_run_async_fallback_can_target_the_requested_group_when_a_pre_router_replaced_it(): + """The requested group was never called when a pre-router selected a tier, so a + tier fallback may legitimately target that originally requested group.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["requested-model"], + original_model_group="requested-model", + original_exception=RuntimeError("selected tier failed"), + max_fallbacks=3, + fallback_depth=0, + model="requested-model", + metadata={"pre_routing_selected_model": "selected-tier"}, + ) + + assert router.received_kwargs["model"] == "requested-model" + assert router.received_kwargs["attempted_targets"].keys == frozenset({"selected-tier", "requested-model"}) + + @pytest.mark.asyncio @pytest.mark.parametrize( "entry", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b413bb18f04..d255135bf64 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8335,20 +8335,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: ) @classmethod - def _router(cls) -> "litellm.Router": + def _router( + cls, + cheap_response: str = "cheap response", + fallbacks: list[dict[str, list[str]]] | None = None, + ) -> "litellm.Router": from litellm.types.router import TaggedPreRoutingStrategy router = litellm.Router( model_list=[ { "model_name": "cheap-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "cheap response"}, + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": cheap_response}, }, { "model_name": "expensive-model", "litellm_params": {"model": "openai/gpt-4o", "mock_response": "expensive response"}, }, - ] + ], + fallbacks=fallbacks, + num_retries=0, ) router.complexity_routers = { "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] @@ -8477,6 +8483,29 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None + @pytest.mark.asyncio + async def test_subagent_can_fallback_to_its_original_requested_model(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"cheap-model": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "expensive response" + assert subagent_kwargs["metadata"]["routing_decision"]["routed_model"] == "cheap-model" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From dba190842cea106ab4b03880861038ddbe6aae42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:26:59 -0700 Subject: [PATCH 18/35] fix(anthropic): keep the cache_control normalizer inside the type-discipline budget --- litellm/llms/anthropic/common_utils.py | 38 +++++++++++++++++--------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index b5bfb32c0c6..19d3d6d7043 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1404,13 +1404,33 @@ def _with_portable_cache_control_in_message(message: object) -> object: return message return { # mutable-ok: JSON wire format **message, - "content": [_with_portable_cache_control_in_content_block(block) for block in content], + "content": [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_content_block(block) for block in content + ], } -def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format +def _with_portable_cache_control_in_messages(messages: object) -> object: + if isinstance(messages, str) or not isinstance(messages, Sequence): + return messages + return [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_message(message) for message in messages + ] + + +def _with_portable_cache_control_in_scoped_value(key: str, value: object) -> object: + match key: + case "system" | "tools": + return _with_portable_cache_control_in_blocks(value) + case "messages": + return _with_portable_cache_control_in_messages(value) + case _: + return value + + +def normalize_cache_control_in_anthropic_payload( payload: Mapping[str, object], -) -> dict[str, object]: +) -> dict[str, object]: # mutable-ok: JSON wire format """ Return a copy of an Anthropic /v1/messages payload with every ``cache_control`` entry reduced to ``{"type": }`` @@ -1427,17 +1447,9 @@ def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire forma dropped entirely. The caller's payload is never mutated. """ portable: Final = _with_portable_cache_control(payload) - scoped: Final = { # mutable-ok: JSON wire format - key: ( - _with_portable_cache_control_in_blocks(value) - if key in ("system", "tools") - else [_with_portable_cache_control_in_message(message) for message in value] - if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) - else value - ) - for key, value in portable.items() + return { # mutable-ok: JSON wire format + key: _with_portable_cache_control_in_scoped_value(key, value) for key, value in portable.items() } - return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: From dbc126cfc97734e47066ab988c3ac1090e7f830a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:07:30 -0700 Subject: [PATCH 19/35] fix(hosted_vllm): forward truncate_prompt_tokens on rerank requests --- .../llms/hosted_vllm/rerank/transformation.py | 21 ++- litellm/types/rerank.py | 21 ++- .../test_hosted_vllm_rerank_transformation.py | 122 +++++++++++++++++- 3 files changed, 154 insertions(+), 10 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0e8fa294f5d..265eb350fc6 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -3,6 +3,7 @@ Transformation logic for Hosted VLLM rerank """ from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -13,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + HostedVLLMRerankTruncationParams, OptionalRerankParams, RerankBilledUnits, RerankRequest, @@ -62,7 +64,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): "top_n", "rank_fields", "return_documents", + "max_tokens_per_doc", "instruction", + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", ] def map_cohere_rerank_params( @@ -100,7 +106,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - return dict(mapped_params) + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + forwarded: Final[OptionalRerankParams] = { + **mapped_params, + "max_tokens_per_doc": max_tokens_per_doc, + "truncate_prompt_tokens": truncation.truncate_prompt_tokens, + "truncation_side": truncation.truncation_side, + "max_tokens_per_query": truncation.max_tokens_per_query, + } + return dict(forwarded) def validate_environment( self, @@ -138,6 +152,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if "documents" not in optional_rerank_params: raise ValueError("documents is required for Hosted VLLM rerank") + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params) rerank_request: Final = RerankRequest( model=model, query=optional_rerank_params["query"], @@ -146,6 +161,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig): rank_fields=optional_rerank_params.get("rank_fields", None), return_documents=optional_rerank_params.get("return_documents", None), instruction=optional_rerank_params.get("instruction", None), + max_tokens_per_doc=truncation.max_tokens_per_doc, + truncate_prompt_tokens=truncation.truncate_prompt_tokens, + truncation_side=truncation.truncation_side, + max_tokens_per_query=truncation.max_tokens_per_query, ) return rerank_request.model_dump(exclude_none=True) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 903781b2ccd..a76e6cf1187 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from pydantic import BaseModel, PrivateAttr -from typing_extensions import Required, TypedDict +from typing import Literal + +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing_extensions import ReadOnly, Required, TypedDict class RerankRequest(BaseModel): @@ -21,6 +23,18 @@ class RerankRequest(BaseModel): # (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing # request when None, so this is fully backward-compatible. instruction: str | None = None + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + + +class HostedVLLMRerankTruncationParams(BaseModel): + model_config = ConfigDict(frozen=True) + + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + max_tokens_per_doc: int | None = None class OptionalRerankParams(TypedDict, total=False): @@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate_prompt_tokens: ReadOnly[int | None] + truncation_side: ReadOnly[Literal["left", "right"] | None] + max_tokens_per_query: ReadOnly[int | None] class RerankBilledUnits(TypedDict, total=False): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index e6e6aa946d5..da27ea1ac58 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,13 @@ +import json import os import sys +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( @@ -87,9 +92,7 @@ class TestHostedVLLMRerankTransform: assert "instruction" not in body def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises( - ValueError, match="Hosted VLLM does not support max_chunks_per_doc" - ): + with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, @@ -104,12 +107,10 @@ class TestHostedVLLMRerankTransform: url = self.config.get_complete_url(base, self.model) assert url == "https://api.example.com/rerank" # Already ends with /rerank - url2 = self.config.get_complete_url( - "https://api.example.com/rerank", self.model - ) + url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): + with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"): self.config.get_complete_url(None, self.model) def test_transform_response(self): @@ -173,3 +174,110 @@ class TestGetOptionalRerankParamsInstruction: documents=["doc1", "doc2"], ) assert "instruction" not in params + + +class TestHostedVLLMRerankTruncationParams: + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): + params = self.config.map_cohere_rerank_params( + non_default_params={ + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "metadata": {"user_api_key": "sk-test"}, + }, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + max_tokens_per_doc=128, + ) + assert params["truncate_prompt_tokens"] == 512 + assert params["truncation_side"] == "left" + assert params["max_tokens_per_query"] == 64 + assert params["max_tokens_per_doc"] == 128 + assert "metadata" not in params + + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): + params = self.config.map_cohere_rerank_params( + non_default_params={"metadata": {"user_api_key": "sk-test"}}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + assert not truncation_keys & body.keys() + assert body == { + "model": self.model, + "query": "test query", + "documents": ["doc1", "doc2"], + "return_documents": True, + } + + def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): + with pytest.raises(ValueError, match="truncation_side"): + self.config.map_cohere_rerank_params( + non_default_params={"truncation_side": "middle"}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + + def test_transform_request_forwards_truncation_params(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": "test query", + "documents": ["doc1", "doc2"], + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "max_tokens_per_doc": 128, + }, + headers={}, + ) + assert body["truncate_prompt_tokens"] == 512 + assert body["truncation_side"] == "left" + assert body["max_tokens_per_query"] == 64 + assert body["max_tokens_per_doc"] == 128 + + def test_transform_request_omits_truncation_params_when_absent(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, + headers={}, + ) + assert "truncate_prompt_tokens" not in body + assert "truncation_side" not in body + assert "max_tokens_per_query" not in body + assert "max_tokens_per_doc" not in body + + def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): + client = HTTPHandler() + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "score-1", + "results": [{"index": 0, "relevance_score": 0.5}], + "usage": {"total_tokens": 512}, + } + with patch.object(client, "post", return_value=mock_response) as mock_post: + litellm.rerank( + model="hosted_vllm/BAAI/bge-reranker-base", + api_base="http://vllm.local:8000", + query="List all the unique case ids", + documents=["a document longer than the reranker context window"], + truncate_prompt_tokens=512, + truncation_side="left", + client=client, + ) + sent_body = json.loads(mock_post.call_args.kwargs["data"]) + assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" + assert sent_body["truncate_prompt_tokens"] == 512 + assert sent_body["truncation_side"] == "left" From 8d00220acef335234797b00cf3b1f1ee73db2b3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:19:20 -0700 Subject: [PATCH 20/35] test(hosted_vllm): annotate rerank truncation test locals as Final --- .../test_hosted_vllm_rerank_transformation.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index da27ea1ac58..49d58cc28c4 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,6 +1,7 @@ import json import os import sys +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -182,7 +183,7 @@ class TestHostedVLLMRerankTruncationParams: self.model = "hosted-vllm-model" def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={ "truncate_prompt_tokens": 512, "truncation_side": "left", @@ -202,15 +203,20 @@ class TestHostedVLLMRerankTruncationParams: assert "metadata" not in params def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], ) - body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) - truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys: Final = { + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", + "max_tokens_per_doc", + } assert not truncation_keys & body.keys() assert body == { "model": self.model, @@ -230,7 +236,7 @@ class TestHostedVLLMRerankTruncationParams: ) def test_transform_request_forwards_truncation_params(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={ "query": "test query", @@ -248,7 +254,7 @@ class TestHostedVLLMRerankTruncationParams: assert body["max_tokens_per_doc"] == 128 def test_transform_request_omits_truncation_params_when_absent(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, headers={}, @@ -259,8 +265,8 @@ class TestHostedVLLMRerankTruncationParams: assert "max_tokens_per_doc" not in body def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): - client = HTTPHandler() - mock_response = MagicMock(spec=httpx.Response) + client: Final = HTTPHandler() + mock_response: Final = MagicMock(spec=httpx.Response) mock_response.status_code = 200 mock_response.json.return_value = { "id": "score-1", @@ -277,7 +283,7 @@ class TestHostedVLLMRerankTruncationParams: truncation_side="left", client=client, ) - sent_body = json.loads(mock_post.call_args.kwargs["data"]) + sent_body: Final = json.loads(mock_post.call_args.kwargs["data"]) assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" assert sent_body["truncate_prompt_tokens"] == 512 assert sent_body["truncation_side"] == "left" From ef14bed0296b4a6c48777b8f9df8018b11179c4b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:31:19 -0700 Subject: [PATCH 21/35] fix(hosted_vllm): reject invalid rerank truncation params with a 400 --- .../llms/hosted_vllm/rerank/transformation.py | 11 +++++++- .../test_hosted_vllm_rerank_transformation.py | 26 ++++++++++++------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 265eb350fc6..764d80c6f82 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -7,8 +7,10 @@ from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError from litellm._uuid import uuid +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -36,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams: + try: + return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + except ValidationError as error: + raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error + + class HostedVLLMRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass @@ -106,7 +115,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + truncation: Final = validated_truncation_params(non_default_params) forwarded: Final[OptionalRerankParams] = { **mapped_params, "max_tokens_per_doc": max_tokens_per_doc, diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 49d58cc28c4..9a62fcf6f0f 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -202,6 +202,22 @@ class TestHostedVLLMRerankTruncationParams: assert params["max_tokens_per_doc"] == 128 assert "metadata" not in params + @pytest.mark.parametrize( + "bad_params", + [{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}], + ) + def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]): + with pytest.raises(litellm.UnsupportedParamsError) as raised: + self.config.map_cohere_rerank_params( + non_default_params=dict(bad_params), + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert raised.value.status_code == 400 + assert next(iter(bad_params)) in str(raised.value) + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, @@ -225,16 +241,6 @@ class TestHostedVLLMRerankTruncationParams: "return_documents": True, } - def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): - with pytest.raises(ValueError, match="truncation_side"): - self.config.map_cohere_rerank_params( - non_default_params={"truncation_side": "middle"}, - model=self.model, - drop_params=False, - query="test query", - documents=["doc1", "doc2"], - ) - def test_transform_request_forwards_truncation_params(self): body: Final = self.config.transform_rerank_request( model=self.model, From d7ee215c57af44219d5a19f5042343ce98e9de2a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:32:04 -0700 Subject: [PATCH 22/35] fix(responses): keep namespace tools intact when a guardrail returns them unchanged Any pre_call guardrail on /v1/responses flattened Codex namespace tools into ns__member functions and wrote the flattened list back to the request, so the model called mcp__server__tool with no namespace and Codex rejected the call as unsupported. The handler now keeps the client's original tools, hands the guardrail a deep copy of the flattened ones, and rebuilds data["tools"] by matching the guardrail's output to the originals by type and name. Unchanged tools go back as the original objects, a dropped or edited namespace member changes only that member, and tools the guardrail injects are still appended. Fixes #39183 --- basedpyright-code-budget.json | 12 +- .../guardrail_translation/handler.py | 114 +++------- .../guardrail_translation/tool_merge.py | 177 +++++++++++++++ .../transformation.py | 161 ++++++++------ ruff-strict-budget.json | 2 +- ...test_openai_responses_guardrail_handler.py | 206 +++++++++++++++++- ...t_openai_responses_guardrail_tool_merge.py | 144 ++++++++++++ .../test_litellm_completion_responses.py | 13 ++ type-discipline-budget.json | 10 +- 9 files changed, 669 insertions(+), 170 deletions(-) create mode 100644 litellm/llms/openai/responses/guardrail_translation/tool_merge.py create mode 100644 tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..e7a069de29a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14076 }, "reportArgumentType": { - "limit": 2216 + "limit": 2215 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4128 + "limit": 4127 }, "reportFunctionMemberAccess": { "limit": 7 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44364 + "limit": 44362 }, "reportUnknownLambdaType": { "limit": 109 @@ -117,13 +117,13 @@ "limit": 111 }, "reportUnnecessaryComparison": { - "limit": 692 + "limit": 687 }, "reportUnnecessaryContains": { - "limit": 5 + "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..5a5970fb867 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,6 +28,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ +import copy import time import uuid from collections.abc import Mapping, Sequence @@ -36,7 +37,6 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from openai.types.responses.tool_param import FunctionToolParam from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -49,6 +49,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -62,7 +63,6 @@ from litellm.types.llms.openai import ( ContentPartDonePartOutputText, ErrorEvent, ErrorEventError, - OpenAIMcpServerTool, OutputItemAddedEvent, OutputItemDoneEvent, OutputTextDeltaEvent, @@ -157,23 +157,31 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - tools_to_check: Final[list[ChatCompletionToolParam]] = [] if input_data is None: return data structured_messages: Final = self.get_structured_messages(data) + raw_tools: Final = data.get("tools") + original_tools: Final[tuple[Mapping[str, object], ...]] = ( + tuple(raw_tools) if isinstance(raw_tools, list) else () + ) + flattened_tool_groups: Final = tuple( + form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) + ) + flattened_tools: Final = tuple( + cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list + for group in flattened_tool_groups + for tool in group + ) + tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list + copy.deepcopy(flattened_tools) + ) # Handle simple string input if isinstance(input_data, str): inputs = GenericGuardrailAPIInputs(texts=[input_data]) - original_tools: list[dict[str, object]] = [] - - # Extract and transform tools if present - if "tools" in data and data["tools"]: - original_tools = list(data["tools"]) - self._extract_and_transform_tools(data["tools"], tools_to_check) - if tools_to_check: - inputs["tools"] = tools_to_check + if tools_to_check: + inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # Include model information if available @@ -189,7 +197,9 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -200,7 +210,6 @@ class OpenAIResponsesHandler(BaseTranslation): texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or []) # Step 1: Extract all text content, images, and tools for msg_idx, message in enumerate(input_data): @@ -212,10 +221,6 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) - # Extract and transform tools if present - if "tools" in data and data["tools"]: - self._extract_and_transform_tools(data["tools"], tools_to_check) - # Step 2: Apply guardrail to all texts in batch if texts_to_check: inputs = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -238,9 +243,7 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_texts = guardrailed_inputs.get("texts", []) self._apply_guardrailed_tools_to_data( - data, - original_tools_list, - guardrailed_inputs.get("tools"), + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) # Step 3: Map guardrail responses back to original input structure @@ -267,73 +270,18 @@ class OpenAIResponsesHandler(BaseTranslation): names.append(str(tool["server_label"])) return names - def _extract_and_transform_tools( - self, - tools: list[FunctionToolParam | OpenAIMcpServerTool], - tools_to_check: list[ChatCompletionToolParam], - ) -> None: - """ - Extract and transform tools from Responses API format to Chat Completion format. - - Uses the LiteLLM transformation function to convert Responses API tools - to Chat Completion tools that can be passed to guardrails. - """ - if tools is not None and isinstance(tools, list): - # Transform Responses API tools to Chat Completion tools - ( - transformed_tools, - _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools) - tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools)) - - def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]: - """ - Remap guardrail-returned tools (Chat Completion format) back to - Responses API request tool format. - """ - return LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( - guardrailed_tools - ) - - def _merge_tools_after_guardrail( - self, - original_tools: list[dict[str, object]], - remapped: list[dict[str, object]], - ) -> list[dict[str, object]]: - """ - Merge remapped guardrailed tools with original tools that were not sent - to the guardrail (e.g. web_search, web_search_preview), preserving order. - Tools a guardrail appended (``remapped`` longer than ``original_tools``) - have no original slot and are kept so an injected tool is not dropped. - """ - if not original_tools: - return remapped - result: Final[list[dict[str, object]]] = [] - j = 0 - for tool in original_tools: - if isinstance(tool, dict) and tool.get("type") in ( - "web_search", - "web_search_preview", - ): - result.append(tool) - else: - if j < len(remapped): - result.append(remapped[j]) - j += 1 - # Keep guardrail-appended tools that matched no original slot above. - result.extend(remapped[j:]) - return result - def _apply_guardrailed_tools_to_data( self, data: dict, - original_tools: list[dict[str, object]], - guardrailed_tools: list[ChatCompletionToolParam] | None, + original_tools: Sequence[Mapping[str, object]], + flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]], + guardrailed_tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - """Remap guardrailed tools to Responses API format and merge with original, then set data['tools'].""" - if guardrailed_tools is not None: - remapped: Final = self._remap_tools_to_responses_api_format(guardrailed_tools) - data["tools"] = self._merge_tools_after_guardrail(original_tools, remapped) + if guardrailed_tools is None: + return + data["tools"] = list( # mutable-ok: downstream wants a list # rebind-ok: in-place request rewrite + merge_guardrailed_tools(original_tools, flattened_tool_groups, guardrailed_tools) + ) def _extract_input_text_and_images( self, diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py new file mode 100644 index 00000000000..3ae951d3f61 --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -0,0 +1,177 @@ +from collections.abc import Iterable, Mapping, Sequence +from itertools import accumulate, chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.responses.litellm_completion_transformation.transformation import ( + NAMESPACE_DESCRIPTION_SEPARATOR, + LiteLLMCompletionResponsesConfig, +) + +Tool: TypeAlias = Mapping[str, object] +IndexedKey: TypeAlias = tuple[str, int] + +_TOOL_ADAPTER: Final = TypeAdapter(dict[str, object]) +_CHAT_TOOL_TOP_LEVEL_KEYS: Final = frozenset({"type", "function"}) + + +def _as_tool(value: object) -> Tool | None: + try: + return _TOOL_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]: + validated: Final = tuple(map(_as_tool, values)) + dropped: Final = sum(tool is None for tool in validated) + if dropped: + verbose_logger.warning("Dropping %d guardrail-returned tools that are not objects", dropped) + return tuple(tool for tool in validated if tool is not None) + + +def _is_function(tool: Tool) -> bool: + return tool.get("type") == "function" + + +def _chat_tool_key(tool: Tool) -> str: + tool_type: Final = str(tool.get("type") or "") + function: Final = _as_tool(tool.get("function")) + if function is not None: + return f"{tool_type}:{function.get('name') or ''}" + return f"{tool_type}:{tool.get('server_label') or tool.get('name') or ''}" + + +def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]: + keys: Final = tuple(_chat_tool_key(tool) for tool in tools) + return tuple((key, keys[:position].count(key)) for position, key in enumerate(keys)) + + +def _namespace_members(namespace: Tool) -> tuple[Tool, ...]: + members: Final = namespace.get("tools") + if not isinstance(members, Sequence) or isinstance(members, (str, bytes)): + return () + return tuple(member for member in map(_as_tool, members) if member is not None) + + +def _function_fields(tool: Tool) -> Tool: + function: Final = _as_tool(tool.get("function")) + return function if function is not None else MappingProxyType({}) + + +def _without_namespace_prefix(key: str, value: object, prefix: str) -> object: + if key != "description" or not isinstance(value, str) or not value.startswith(prefix): + return value + return value[len(prefix) :] + + +def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: + flattened_function: Final = _function_fields(flattened) + prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" + changed_function: Final = MappingProxyType( + { + key: _without_namespace_prefix(key, value, prefix) + for key, value in _function_fields(guardrailed).items() + if flattened_function.get(key) != value + } + ) + changed_extras: Final = MappingProxyType( + { + key: value + for key, value in guardrailed.items() + if key not in _CHAT_TOOL_TOP_LEVEL_KEYS and flattened.get(key) != value + } + ) + return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType + + +def _rebuilt_function_members( + function_members: Sequence[Tool], + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], + namespace_description: str, +) -> tuple[Tool | None, ...]: + return tuple( + None + if key not in guardrailed_by_key + else member + if guardrailed_by_key[key] == flattened + else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description) + for member, flattened, key in zip(function_members, flattened_group, group_keys) + ) + + +def _rebuilt_namespace( + original: Tool, + members: Sequence[Tool], + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], +) -> tuple[Tool, ...]: + namespace_description: Final = str(original.get("description") or "") + rebuilt_functions: Final = iter( + _rebuilt_function_members( + tuple(member for member in members if _is_function(member)), + flattened_group, + group_keys, + guardrailed_by_key, + namespace_description, + ) + ) + rebuilt_members: Final = tuple( + rebuilt + for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members) + if rebuilt is not None + ) + if not rebuilt_members: + return () + return ({**original, "tools": list(rebuilt_members)},) # mutable-ok: json.dumps needs a plain dict and list + + +def _merged_original( + original: Tool, + flattened_group: Sequence[Tool], + group_keys: Sequence[IndexedKey], + guardrailed_by_key: Mapping[IndexedKey, Tool], +) -> tuple[Tool, ...]: + if not group_keys: + return (original,) + guardrailed_group: Final = tuple(guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key) + if guardrailed_group == tuple(flattened_group): + return (original,) + if not guardrailed_group: + return () + members: Final = _namespace_members(original) if original.get("type") == "namespace" else () + if members and sum(map(_is_function, members)) == len(flattened_group): + return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) + return tuple( + LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(guardrailed_group) + ) + + +def merge_guardrailed_tools( + original_tools: Sequence[Tool], + flattened_groups: Sequence[Sequence[Tool]], + guardrailed_tools: Iterable[object], +) -> tuple[Tool, ...]: + guardrailed: Final = _validated_tools(guardrailed_tools) + flattened_keys: Final = _indexed_keys(tuple(chain.from_iterable(flattened_groups))) + guardrailed_keys: Final = _indexed_keys(guardrailed) + guardrailed_by_key: Final = MappingProxyType(dict(zip(guardrailed_keys, guardrailed))) + group_ends: Final = tuple(accumulate(len(group) for group in flattened_groups)) + group_key_slices: Final = tuple( + flattened_keys[end - len(group) : end] for group, end in zip(flattened_groups, group_ends) + ) + merged_originals: Final = chain.from_iterable( + _merged_original(original, group, group_keys, guardrailed_by_key) + for original, group, group_keys in zip(original_tools, flattened_groups, group_key_slices) + ) + owned_keys: Final = frozenset(flattened_keys) + appended: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( + tuple(tool for key, tool in zip(guardrailed_keys, guardrailed) if key not in owned_keys) + ) + return tuple(chain(merged_originals, appended)) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5f3e88bb12f..9f91d7527cb 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -6,6 +6,7 @@ import json import re import uuid from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -102,6 +103,15 @@ from .custom_tools import ( NamespaceNameMap: TypeAlias = Mapping[str, tuple[str, str]] NamespaceTool: TypeAlias = Mapping[str, object] ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None +ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool +NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" + + +@dataclass(frozen=True, slots=True) +class ResponsesToolChatForm: + chat_tools: tuple[ChatToolParam, ...] + web_search_options: OpenAIWebSearchOptions | None + if TYPE_CHECKING: from openai.types.responses.response_apply_patch_tool_call import ( @@ -1771,7 +1781,7 @@ class LiteLLMCompletionResponsesConfig: tool_name: Final = str(namespace_tool.get("name") or "") raw_description: Final = str(namespace_tool.get("description") or "") description: Final = ( - f"{namespace_description}\n\n{raw_description}" + f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" if nested and namespace_description and raw_description else namespace_description if nested and namespace_description @@ -1837,9 +1847,82 @@ class LiteLLMCompletionResponsesConfig: + ", ".join(sorted(conflicting_tool_names)) ) + @staticmethod + def _responses_tool_to_chat_form(tool: Mapping[str, object]) -> ResponsesToolChatForm: + tool_type: Final = tool.get("type") + if tool_type == "mcp": + return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) + if tool_type == "web_search_preview" or tool_type == "web_search": + _search_context_size: Final[Literal["low", "medium", "high"]] = cast( + Literal["low", "medium", "high"], tool.get("search_context_size") + ) + _user_location: Final[OpenAIWebSearchUserLocation | None] = cast( + OpenAIWebSearchUserLocation | None, + tool.get("user_location") or None, + ) + return ResponsesToolChatForm( + chat_tools=(), + web_search_options=OpenAIWebSearchOptions( + search_context_size=_search_context_size, + user_location=_user_location, + ), + ) + if tool_type == "function": + typed_tool: Final = cast(FunctionToolParam, tool) + raw_parameters: Final = typed_tool.get("parameters", {}) or {} + parameters: Final = ( + {**raw_parameters} # mutable-ok: json.dumps rejects MappingProxyType + if "type" in raw_parameters + else {**raw_parameters, "type": "object"} # mutable-ok: json.dumps rejects MappingProxyType + ) + chat_completion_tool: Final[dict[str, object]] = { + "type": "function", + "function": { + "name": typed_tool.get("name") or "", + "description": typed_tool.get("description") or "", + "parameters": parameters, + "strict": typed_tool.get("strict", False) or False, + }, + } + if tool.get("cache_control"): + chat_completion_tool["cache_control"] = tool.get("cache_control") + if tool.get("defer_loading"): + chat_completion_tool["defer_loading"] = tool.get("defer_loading") + if tool.get("allowed_callers"): + chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") + if tool.get("input_examples"): + chat_completion_tool["input_examples"] = tool.get("input_examples") + return ResponsesToolChatForm( + chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None + ) + if tool_type == "namespace": + return ResponsesToolChatForm( + chat_tools=LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool), web_search_options=None + ) + if tool_type == "custom": + converted: Final = convert_custom_tool_to_function_tool(tool) + return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) + if tool_type in ("computer_use", "image_generation", "shell"): + # Drop unsupported Responses-API-only tool types that have no + # Chat Completions equivalent. Passing them through verbatim + # causes providers to reject the request with "'function' is a + # required property". + verbose_logger.warning( + "Dropping Responses API tool of type '%s': it has no Chat Completions " + "equivalent and the target provider would reject the request.", + tool_type, + ) + return ResponsesToolChatForm(chat_tools=(), web_search_options=None) + return ResponsesToolChatForm(chat_tools=(cast(ChatToolParam, tool),), web_search_options=None) + + @staticmethod + def responses_tools_to_chat_forms(tools: ResponseTools) -> tuple[ResponsesToolChatForm, ...]: + LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) + return tuple(LiteLLMCompletionResponsesConfig._responses_tool_to_chat_form(tool) for tool in tools or ()) + @staticmethod def transform_responses_api_tools_to_chat_completion_tools( - tools: list[FunctionToolParam | OpenAIMcpServerTool] | None, + tools: ResponseTools, ) -> tuple[ list[ChatCompletionToolParam | OpenAIMcpServerTool], OpenAIWebSearchOptions | None, @@ -1849,73 +1932,16 @@ class LiteLLMCompletionResponsesConfig: """ if tools is None: return [], None - LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) - chat_completion_tools: Final[list[ChatCompletionToolParam | OpenAIMcpServerTool]] = [] - web_search_options: OpenAIWebSearchOptions | None = None - for tool in tools: - if tool.get("type") == "mcp": - chat_completion_tools.append(cast(OpenAIMcpServerTool, tool)) - elif tool.get("type") == "web_search_preview" or tool.get("type") == "web_search": - _search_context_size: Literal["low", "medium", "high"] = cast( - Literal["low", "medium", "high"], tool.get("search_context_size") - ) - _user_location: OpenAIWebSearchUserLocation | None = cast( - OpenAIWebSearchUserLocation | None, - tool.get("user_location") or None, - ) - web_search_options = OpenAIWebSearchOptions( - search_context_size=_search_context_size, - user_location=_user_location, - ) - elif tool.get("type") == "function": - typed_tool = cast(FunctionToolParam, tool) - # Ensure parameters has "type": "object" as required by providers like Anthropic - parameters = dict(typed_tool.get("parameters", {}) or {}) - if not parameters or "type" not in parameters: - parameters["type"] = "object" - chat_completion_tool: dict[str, object] = { - "type": "function", - "function": { - "name": typed_tool.get("name") or "", - "description": typed_tool.get("description") or "", - "parameters": parameters, - "strict": typed_tool.get("strict", False) or False, - }, - } - if tool.get("cache_control"): - chat_completion_tool["cache_control"] = tool.get("cache_control") - if tool.get("defer_loading"): - chat_completion_tool["defer_loading"] = tool.get("defer_loading") - if tool.get("allowed_callers"): - chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") - if tool.get("input_examples"): - chat_completion_tool["input_examples"] = tool.get("input_examples") - chat_completion_tools.append(cast(ChatCompletionToolParam, chat_completion_tool)) - elif tool.get("type") == "namespace": - chat_completion_tools.extend(LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool)) - elif tool.get("type") == "custom": - converted = convert_custom_tool_to_function_tool(tool) - if converted is not None: - chat_completion_tools.append(converted) - else: - _tool_type = tool.get("type") - if _tool_type in ("computer_use", "image_generation", "shell"): - # Drop unsupported Responses-API-only tool types that have no - # Chat Completions equivalent. Passing them through verbatim - # causes providers to reject the request with "'function' is a - # required property". - verbose_logger.warning( - "Dropping Responses API tool of type '%s': it has no Chat Completions " - "equivalent and the target provider would reject the request.", - _tool_type, - ) - continue - chat_completion_tools.append(cast(ChatCompletionToolParam | OpenAIMcpServerTool, tool)) - return chat_completion_tools, web_search_options + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools) + web_search_options: Final = next( + (form.web_search_options for form in reversed(forms) if form.web_search_options is not None), + None, + ) + return [chat_tool for form in forms for chat_tool in form.chat_tools], web_search_options @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( - chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, + chat_completion_tools: Sequence[Mapping[str, object]] | None, ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to @@ -1926,9 +1952,6 @@ class LiteLLMCompletionResponsesConfig: return [] result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: - if not isinstance(tool, dict): - result.append(tool) - continue if tool.get("type") == "function": fn = cast(_ToolFunctionDefinition, tool.get("function") or {}) parameters = dict(fn.get("parameters", {}) or {}) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..935fed18a79 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 31 + "limit": 29 }, "RUF046": { "limit": 4 diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 315b6948bd8..d071ef78c2d 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -5,6 +5,8 @@ Tests the handler's ability to process input/output for the Responses API with guardrail transformations. """ +import copy +from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock @@ -19,6 +21,10 @@ from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.responses.main import GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -1287,14 +1293,14 @@ class TestOpenAIResponsesHandlerToolInjection: """A tool a guardrail injects must survive the write-back to Responses format.""" def test_merge_keeps_guardrail_appended_tool(self): - """_merge_tools_after_guardrail must not drop the extra appended tool.""" - handler = OpenAIResponsesHandler() + """merge_guardrailed_tools must not drop the extra appended tool.""" original = [{"type": "function", "name": "a"}] - remapped = [ - {"type": "function", "name": "a"}, - {"type": "function", "name": "b"}, + groups = [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original)] + guardrailed = [ + *groups[0], + {"type": "function", "function": {"name": "b", "description": "", "parameters": {"type": "object"}}}, ] - merged = handler._merge_tools_after_guardrail(original, remapped) + merged = merge_guardrailed_tools(original, groups, guardrailed) assert [t["name"] for t in merged] == ["a", "b"] @pytest.mark.asyncio @@ -1323,6 +1329,194 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +class ToolEditingGuardrail(CustomGuardrail): + """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" + + def __init__(self, edit: Callable[[list[dict]], list[dict]], **kwargs): + super().__init__(**kwargs) + self.edit = edit + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = self.edit(list(inputs.get("tools") or [])) + return inputs + + +def _codex_request(input_value): + """A Responses API request shaped like what the Codex CLI sends when an MCP server is configured""" + return { + "model": "gpt-5.3-codex", + "input": input_value, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Weather lookup", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "strict": False, + }, + { + "type": "namespace", + "name": "mcp__confluence", + "description": "Confluence tools", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + "strict": False, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search pages", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + "strict": False, + }, + ], + }, + { + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + "format": {"type": "grammar", "syntax": "lark", "definition": 'start: "x"'}, + }, + {"type": "web_search"}, + ], + } + + +def _tool_named(tools, name): + return next(tool for tool in tools if tool.get("name") == name) + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Codex sends MCP tools as ``namespace`` tools; a guardrail must never flatten them (GH #39183)""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "input_value", + ["hi", [{"role": "user", "content": "hi", "type": "message"}]], + ids=["string_input", "list_input"], + ) + async def test_pass_through_guardrail_leaves_tools_untouched(self, input_value): + data = _codex_request(input_value) + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, MockPassThroughGuardrail(guardrail_name="test") + ) + + assert result["tools"] == expected_tools + + @pytest.mark.asyncio + async def test_appending_guardrail_keeps_namespace_and_adds_tool(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolAppendingGuardrail(guardrail_name="test") + ) + + assert result["tools"][:-1] == expected_tools + assert result["tools"][-1]["type"] == "function" + assert result["tools"][-1]["name"] == "injected_tool" + + @pytest.mark.asyncio + async def test_dropping_one_member_prunes_only_that_member(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if t["function"]["name"] != "mcp__confluence__confluence_search"], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert [member["name"] for member in namespace["tools"]] == ["confluence_get_page"] + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert [t for t in result["tools"] if t is not namespace] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_editing_a_member_lands_on_that_member_without_the_namespace_prefix(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def redact_search(tools): + for tool in tools: + if tool["function"]["name"] == "mcp__confluence__confluence_search": + tool["function"]["description"] = "Confluence tools\n\nREDACTED" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=redact_search, guardrail_name="test") + ) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert namespace["tools"][1] == {**expected_tools[1]["tools"][1], "description": "REDACTED"} + assert {k: v for k, v in namespace.items() if k != "tools"} == { + k: v for k, v in expected_tools[1].items() if k != "tools" + } + + @pytest.mark.asyncio + async def test_dropping_every_member_drops_the_namespace(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if not t["function"]["name"].startswith("mcp__confluence__")], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["tools"] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_edited_top_level_function_is_rewritten_in_place(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def rename_weather(tools): + for tool in tools: + if tool["function"]["name"] == "get_weather": + tool["function"]["description"] = "Weather lookup (guarded)" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=rename_weather, guardrail_name="test") + ) + + assert result["tools"][0] == {**expected_tools[0], "description": "Weather lookup (guarded)"} + assert result["tools"][1:] == expected_tools[1:] + + +class TestOpenAIResponsesHandlerMalformedTools: + @pytest.mark.asyncio + async def test_request_tools_that_are_not_a_list_never_reach_the_guardrail(self): + handler = OpenAIResponsesHandler() + seen: list[list[dict]] = [] + + def record(tools): + seen.append(tools) + return tools + + guardrail = ToolEditingGuardrail(edit=record, guardrail_name="test") + data = {"input": "hi", "tools": {"type": "function", "name": "get_weather"}} + + result = await handler.process_input_messages(data, guardrail) + + assert seen == [[]] + assert result["input"] == "hi" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py new file mode 100644 index 00000000000..b80dd0b36aa --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -0,0 +1,144 @@ +""" +Unit tests for merge_guardrailed_tools, which writes guardrail-returned chat tools back onto the +Responses API request tools they were flattened from +""" + +import copy + +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +def _groups(tools): + return [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools)] + + +def _flat(groups): + return [chat_tool for group in groups for chat_tool in group] + + +def _function(name, description=""): + return {"type": "function", "name": name, "description": description, "parameters": {"type": "object"}} + + +def test_unchanged_tools_come_back_as_the_original_objects(): + original = [ + _function("a"), + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x"), _function("y")]}, + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "web_search"}, + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_guardrail_reordering_unchanged_tools_keeps_request_order(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}, {"type": "web_search"}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, list(reversed(_flat(groups)))) + + assert list(merged) == original + + +def test_duplicate_function_names_are_matched_by_ordinal(): + original = [_function("dup", "first"), _function("dup", "second")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)[:1]) + + assert list(merged) == [original[0]] + + +def test_edited_mcp_tool_is_rewritten(): + original = [{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}] + groups = _groups(original) + edited = [{**groups[0][0], "allowed_tools": ["read_wiki_structure"]}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == edited + + +def test_injected_tool_lands_after_the_request_tools_when_request_had_none(): + injected = {"type": "function", "function": {"name": "b", "description": "d", "parameters": {"type": "object"}}} + + merged = merge_guardrailed_tools([], [], [injected]) + + assert list(merged) == [ + {"type": "function", "name": "b", "description": "d", "parameters": {"type": "object"}, "strict": False} + ] + + +def test_empty_guardrail_output_keeps_only_tools_never_sent_to_the_guardrail(): + original = [_function("a"), {"type": "web_search"}, {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + + merged = merge_guardrailed_tools(original, _groups(original), []) + + assert list(merged) == [{"type": "web_search"}] + + +def test_member_edit_strips_only_the_namespace_description_prefix(): + original = [{"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + assert groups[0][0]["function"]["description"] == "NS\n\nX doc" + edited = [{**groups[0][0], "function": {**groups[0][0]["function"], "description": "NS\n\nX doc (guarded)"}}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc (guarded)")]} + ] + + +def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read", "Read"), custom_member]} + ] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "NS\n\nEDITED" + + merged = merge_guardrailed_tools(original, groups, edited) + + assert len(merged) == 1 + assert [member["name"] for member in merged[0]["tools"]] == ["read", "grep"] + assert merged[0]["tools"][0]["description"] == "EDITED" + assert merged[0]["tools"][1] == custom_member + + +def test_member_extras_edited_by_the_guardrail_land_on_that_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["cache_control"] = {"type": "ephemeral"} + + merged = merge_guardrailed_tools(original, groups, edited) + + assert merged[0]["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert merged[0]["tools"][0]["name"] == "read" + + +def test_guardrail_output_is_read_once(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, (chat_tool for chat_tool in _flat(groups))) + + assert list(merged) == original + + +def test_non_object_guardrail_items_are_dropped(): + original = [_function("a")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [*_flat(groups), "junk", None]) + + assert list(merged) == original diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b2b8eb5da80..2068f10ea2d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1928,6 +1928,19 @@ class TestToolTransformation: assert result_tool["function"]["parameters"]["type"] == "object" assert "properties" in result_tool["function"]["parameters"] + def test_transform_function_tools_parameters_keep_client_key_order(self): + tools = [ + {"type": "function", "name": "a", "parameters": {"properties": {"arg": {"type": "string"}}, "required": ["arg"]}}, + {"type": "function", "name": "b", "parameters": {"type": "object", "properties": {}}}, + ] + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=tools + ) + + assert list(result_tools[0]["function"]["parameters"]) == ["properties", "required", "type"] + assert list(result_tools[1]["function"]["parameters"]) == ["type", "properties"] + def test_transform_function_tools_empty_parameters(self): """Test that empty parameters get 'type': 'object' added""" function_tool = { diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..970a44cd4fa 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22340 }, "LIT002": { - "limit": 26777 + "limit": 26770 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1038 }, "LIT007": { "limit": 0 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16503 }, "LIT011": { - "limit": 5535 + "limit": 5534 }, "LIT012": { "limit": 4495 From dc12e4c2b4ad31b1eda1544ecd2e424aabc78151 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:57:01 -0700 Subject: [PATCH 23/35] fix(responses): match guardrail tools by ordinal in one pass Sort the chat-tool keys once and number duplicates with groupby instead of rescanning every preceding key per position, so the guardrail merge stays O(n log n) on client-supplied tool lists. Drop the comment that restated the unsupported-tool warning in the Responses-to-chat transformation. --- .../guardrail_translation/tool_merge.py | 8 ++++++-- .../transformation.py | 4 ---- ...st_openai_responses_guardrail_tool_merge.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 3ae951d3f61..9fbf9f31a0f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Mapping, Sequence -from itertools import accumulate, chain +from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias @@ -47,7 +47,11 @@ def _chat_tool_key(tool: Tool) -> str: def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]: keys: Final = tuple(_chat_tool_key(tool) for tool in tools) - return tuple((key, keys[:position].count(key)) for position, key in enumerate(keys)) + positions_by_key: Final = groupby(sorted(range(len(keys)), key=keys.__getitem__), key=keys.__getitem__) + ordinal_by_position: Final = MappingProxyType( + {position: ordinal for _, positions in positions_by_key for ordinal, position in enumerate(positions)} + ) + return tuple((key, ordinal_by_position[position]) for position, key in enumerate(keys)) def _namespace_members(namespace: Tool) -> tuple[Tool, ...]: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 9f91d7527cb..b2d1a69e0d8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1903,10 +1903,6 @@ class LiteLLMCompletionResponsesConfig: converted: Final = convert_custom_tool_to_function_tool(tool) return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) if tool_type in ("computer_use", "image_generation", "shell"): - # Drop unsupported Responses-API-only tool types that have no - # Chat Completions equivalent. Passing them through verbatim - # causes providers to reject the request with "'function' is a - # required property". verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index b80dd0b36aa..4075c209606 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -56,6 +56,24 @@ def test_duplicate_function_names_are_matched_by_ordinal(): assert list(merged) == [original[0]] +def test_interleaved_duplicate_names_keep_their_own_ordinals(): + original = [ + _function("dup", "a"), + _function("other", "x"), + _function("dup", "b"), + _function("dup", "c"), + _function("other", "y"), + ] + groups = _groups(original) + flat = _flat(groups) + edited = {**flat[3], "function": {**flat[3]["function"], "description": "changed"}} + + merged = merge_guardrailed_tools(original, groups, [*flat[:3], edited, flat[4]]) + + assert list(merged) == [*original[:3], {**_function("dup", "changed"), "strict": False}, original[4]] + assert all(merged[position] is original[position] for position in (0, 1, 2, 4)) + + def test_edited_mcp_tool_is_rewritten(): original = [{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}] groups = _groups(original) From f4eca10f1d7f832f6300b588434bcd99003f4bb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:07:57 +0000 Subject: [PATCH 24/35] ci: retrigger checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 0f6d983c7057faf13716639869d828d309bcba5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:08:14 -0700 Subject: [PATCH 25/35] fix(router): skip Claude Code session binding without pre-routing strategies --- litellm/router.py | 2 ++ tests/test_litellm/test_router.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fc4f2d227f5..b1038ca6002 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12636,6 +12636,8 @@ class Router: registered_model_name: str, request_kwargs: Mapping[str, object], ) -> str: + if not any((self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers)): + return registered_model_name cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) if cache_key is None or not isinstance(request_kwargs, dict): return registered_model_name diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 05dcb664ed5..ef25502a4f9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8499,6 +8499,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): + from litellm.caching.caching import RedisCache + + router = self._router() + router.complexity_routers = {} + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock() + redis_cache.async_delete_cache = AsyncMock() + router._update_redis_cache(cache=redis_cache) + + for request_kwargs in (self._request_kwargs(), self._request_kwargs(agent_id="agent-1234")): + response = await router.async_pre_routing_hook(model="expensive-model", request_kwargs=request_kwargs) + assert response is None + + redis_cache.async_get_cache.assert_not_awaited() + redis_cache.async_set_cache.assert_not_awaited() + redis_cache.async_delete_cache.assert_not_awaited() + @pytest.mark.asyncio async def test_session_bindings_do_not_evict_router_rate_limit_state(self): router = self._router() From 4705dc6325fc8b2803e59ad19c85f9d07230a9a8 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:34:54 +0000 Subject: [PATCH 26/35] fix: retry P3009 when the deadlocked ledger row was already rolled back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/utils.py | 46 +++++++---------- .../tests/test_setup_database_fail_fast.py | 51 +++++++++++++++++++ 2 files changed, 70 insertions(+), 27 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 1183757c827..9cc33e6ca56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -277,22 +277,16 @@ class ProxyExtrasDBManager: pass @staticmethod - def _failed_migration_logs(migration_name: str) -> str: - """Logs recorded on the migration's failed _prisma_migrations row. - - P3009 stderr does not carry the original failure, so this is the only - way to tell a migration that lost a deadlock race against a concurrent - migrate deploy from one whose SQL is genuinely broken. Returns "" when - psycopg is missing, the DB is unreachable, or no failed row exists. - """ + def _failed_migration_logs(migration_name: str) -> Optional[str]: + """Return failed migration logs, or None if the ledger is unavailable.""" database_url = os.getenv("DATABASE_URL") if not database_url: - return "" + return None try: import psycopg except ImportError: - return "" + return None cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) ledger_table = psycopg.sql.SQL("{}.{}").format( @@ -314,7 +308,7 @@ class ProxyExtrasDBManager: (migration_name,), ).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): - return "" + return None return (row[0] or "") if row else "" @staticmethod @@ -825,22 +819,20 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue - if migration_match and _MIGRATION_DEADLOCK_MARKER in ( - ProxyExtrasDBManager._failed_migration_logs( - migration_match.group(1) - ) - ): - logger.info( - "Migration %s lost a deadlock race against a " - "concurrent migrate deploy, rolling its ledger " - "row back and retrying", - migration_match.group(1), - ) - ProxyExtrasDBManager._roll_back_migration_best_effort( - migration_match.group(1) - ) - time.sleep(random.randrange(5, 15)) - continue + if migration_match: + migration_name = migration_match.group(1) + ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) + if ledger_logs is not None and ( + ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs + ): + logger.info( + "Migration %s failed in a concurrent migrate deploy " + "deadlock race, rolling its ledger row back and retrying", + migration_name, + ) + ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + time.sleep(random.randrange(5, 15)) + continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index c4347a91dce..70e0f216338 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -350,6 +350,57 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] +def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path): + """v2: empty failed ledger logs mean a concurrent deploy moved it on.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "") + rolled_back = [] + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: rolled_back.append(name), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + lambda name: pytest.fail("a deadlocked migration must never be marked applied"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert rolled_back == ["20260415120000_health_check_latest_per_model_index"] + + +def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): + """v2: an unreadable ledger cannot establish that P3009 was a deadlock.""" + _stub_v2_env(monkeypatch, tmp_path) + + stderr = ( + "Error: P3009\n" + "migrate found failed migrations in the target database\n" + "The `20260415120000_health_check_latest_per_model_index` migration " + "started at 2026-09-01 18:46:13 UTC failed" + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), + ) + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): """v2: a failed ledger row whose logs show a real SQL error stays fatal.""" _stub_v2_env(monkeypatch, tmp_path) From f6eff1bde0f2464ce06c560b55feb7dab4192ab7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:55:16 -0700 Subject: [PATCH 27/35] fix(router): keep Claude Code session bindings across side calls and workers --- litellm/router.py | 9 ++- .../router_code_coverage.py | 1 + tests/test_litellm/test_router.py | 58 +++++++++++++++---- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b1038ca6002..f4d912eb5be 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12630,6 +12630,12 @@ class Router: e, ) + async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: + session_cache: Final = self._claude_code_session_router_cache + if session_cache.redis_cache is None: + return await session_cache.async_get_cache(key=cache_key) + return await session_cache.redis_cache.async_get_cache(key=cache_key) + async def _resolve_claude_code_session_router( self, model: str, @@ -12646,7 +12652,7 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=cache_key) + bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model @@ -12664,7 +12670,6 @@ class Router: if self._request_header(request_kwargs, "x-app") != "cli": return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self._claude_code_session_router_cache.async_set_cache( key=cache_key, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 60b56b7fac6..0af29f069c6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -86,6 +86,7 @@ ignored_function_names = [ "_claude_code_session_router_cache_key", # Tested through Claude Code session routing in test_router.py "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py + "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py ] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ef25502a4f9..1e5767aa645 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8383,18 +8383,21 @@ class TestConsumedRequestTagsStamp: class TestClaudeCodeSubagentSessionRouterBinding: class _RewriteStrategy: + def __init__(self, routed_model: str = "cheap-model") -> None: + self.routed_model = routed_model + async def async_pre_routing_hook( self, model, request_kwargs, messages=None, input=None, specific_deployment=False ): from litellm.types.router import PreRoutingHookResponse return PreRoutingHookResponse( - model="cheap-model", + model=self.routed_model, messages=messages, routing_decision={ "router_model_name": "smart-router", "router_type": "complexity", - "routed_model": "cheap-model", + "routed_model": self.routed_model, "cause": "heuristic_scorer", }, ) @@ -8422,7 +8425,8 @@ class TestClaudeCodeSubagentSessionRouterBinding: num_retries=0, ) router.complexity_routers = { - "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy()),), + "premium-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy("expensive-model")),), } return router @@ -8467,7 +8471,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" @pytest.mark.asyncio - async def test_main_direct_model_clears_the_session_router(self): + async def test_main_thread_side_calls_to_a_plain_model_keep_the_session_router(self): router = self._router() await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) @@ -8478,33 +8482,67 @@ class TestClaudeCodeSubagentSessionRouterBinding: request_kwargs=self._request_kwargs(agent_id="agent-1234"), ) - assert response is None + assert response is not None + assert response.model == "cheap-model" @pytest.mark.asyncio - async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): from litellm.caching.caching import RedisCache router = self._router() + del router.complexity_routers["smart-router"] redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - - await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) router._update_redis_cache(cache=redis_cache) response = await router.async_pre_routing_hook( model="expensive-model", - request_kwargs=self._request_kwargs(), + request_kwargs=self._request_kwargs(agent_id="agent-1234"), ) assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): + from types import SimpleNamespace + + from litellm.caching.caching import RedisCache + + shared_binding = SimpleNamespace(value=None) + shared_redis = MagicMock(spec=RedisCache) + shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) + shared_redis.async_set_cache = AsyncMock( + side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) + ) + main_worker, subagent_worker = self._router(), self._router() + main_worker._update_redis_cache(cache=shared_redis) + subagent_worker._update_redis_cache(cache=shared_redis) + + await main_worker.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + first = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + await main_worker.async_pre_routing_hook(model="premium-router", request_kwargs=self._request_kwargs()) + second = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert first is not None + assert first.model == "cheap-model" + assert second is not None + assert second.model == "expensive-model" + assert shared_binding.value == "premium-router" + @pytest.mark.asyncio async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): from litellm.caching.caching import RedisCache router = self._router() - router.complexity_routers = {} + router.complexity_routers.clear() redis_cache = MagicMock(spec=RedisCache) redis_cache.async_get_cache = AsyncMock(return_value=None) redis_cache.async_set_cache = AsyncMock() From 5561b8438c34054b3fd1475e011ed5cfd6f6ec97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:18:16 -0700 Subject: [PATCH 28/35] fix(responses): keep a namespace's non-function members when every function member is dropped --- .../responses/guardrail_translation/tool_merge.py | 4 ++-- .../test_openai_responses_guardrail_tool_merge.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9fbf9f31a0f..9a34946838a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -147,11 +147,11 @@ def _merged_original( guardrailed_group: Final = tuple(guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key) if guardrailed_group == tuple(flattened_group): return (original,) - if not guardrailed_group: - return () members: Final = _namespace_members(original) if original.get("type") == "namespace" else () if members and sum(map(_is_function, members)) == len(flattened_group): return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) + if not guardrailed_group: + return () return tuple( LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(guardrailed_group) ) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index 4075c209606..cbcddb98235 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -132,6 +132,19 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited( assert merged[0]["tools"][1] == custom_member +def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, + _function("a"), + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [groups[1][0]]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + + def test_member_extras_edited_by_the_guardrail_land_on_that_member(): original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] groups = _groups(original) From 7cbc74399d683b1113a4236768ac8414e7e9f07d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:24:43 -0700 Subject: [PATCH 29/35] fix(responses): accept pydantic tool objects returned by guardrails --- .../guardrail_translation/tool_merge.py | 5 ++-- ...t_openai_responses_guardrail_tool_merge.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9a34946838a..b596adfad6f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -3,7 +3,7 @@ from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.responses.litellm_completion_transformation.transformation import ( @@ -19,8 +19,9 @@ _CHAT_TOOL_TOP_LEVEL_KEYS: Final = frozenset({"type", "function"}) def _as_tool(value: object) -> Tool | None: + candidate: Final = value.model_dump(exclude_unset=True) if isinstance(value, BaseModel) else value try: - return _TOOL_ADAPTER.validate_python(value) + return _TOOL_ADAPTER.validate_python(candidate) except ValidationError: return None diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index cbcddb98235..9c236d81f51 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -9,6 +9,7 @@ from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GuardrailToolParam def _groups(tools): @@ -166,6 +167,28 @@ def test_guardrail_output_is_read_once(): assert list(merged) == original +def test_pydantic_guardrail_tools_round_trip_like_dicts(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + models = [GuardrailToolParam.model_validate(chat_tool) for chat_tool in _flat(groups)] + + merged = merge_guardrailed_tools(original, groups, models) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_pydantic_guardrail_edit_lands_on_the_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "EDITED" + + merged = merge_guardrailed_tools(original, groups, [GuardrailToolParam.model_validate(edited[0])]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "tools": [_function("x", "EDITED")]}] + + def test_non_object_guardrail_items_are_dropped(): original = [_function("a")] groups = _groups(original) From 979c3cb1aa0dc579c38d10f44839b93ecc0d9823 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:29:10 -0700 Subject: [PATCH 30/35] fix(proxy-extras): retry when the advisory lock wait times out behind a concurrent migrate deploy --- .../litellm_proxy_extras/utils.py | 9 ++++++ .../tests/test_setup_database_fail_fast.py | 28 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 9cc33e6ca56..d22484bc0e8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -905,6 +905,15 @@ class ProxyExtrasDBManager: time.sleep(random.randrange(5, 15)) continue + if "P1002" in stderr and "advisory lock" in stderr: + logger.info( + "prisma migrate deploy attempt %s timed out waiting for " + "the advisory lock a concurrent migrate deploy holds, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + continue + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 70e0f216338..406f07eb792 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -432,3 +432,31 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True + + +_P1002_ADVISORY_LOCK_STDERR = ( + "Error: P1002\n\n" + "The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n" + "Context: Timed out trying to acquire a postgres advisory lock " + "(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms." +) + + +def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path): + """v2: the advisory-lock waiter that times out while a peer's retry holds + the lock retries instead of dying.""" + _stub_v2_env(monkeypatch, tmp_path) + monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR)) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + + +def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path): + """v2: a plain P1002 (database unreachable) stays fatal.""" + _stub_v2_env(monkeypatch, tmp_path) + stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out." + monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) From ed0d3f94422e32c498781fe63722219989987285 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 15:21:37 -0700 Subject: [PATCH 31/35] fix(router): fail open on session binding reads --- litellm/router.py | 13 ++++++-- tests/test_litellm/test_router.py | 50 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f4d912eb5be..2d8b4a43759 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12632,9 +12632,16 @@ class Router: async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: session_cache: Final = self._claude_code_session_router_cache - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) + try: + if session_cache.redis_cache is None: + return await session_cache.async_get_cache(key=cache_key) + return await session_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis + verbose_router_logger.warning( + "Failed to read Claude Code session router binding; using the requested model: %s", + e, + ) + return None async def _resolve_claude_code_session_router( self, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1e5767aa645..16c4d997ad7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8504,6 +8504,56 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_redis_read_failure_does_not_reject_a_subagent_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + request_kwargs = self._request_kwargs(agent_id="agent-1234") + cache_key = router._claude_code_session_router_cache_key(request_kwargs) + assert cache_key is not None + await router._claude_code_session_router_cache.in_memory_cache.async_set_cache( + cache_key, + "smart-router", + ) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=request_kwargs, + ) + + assert response is None + assert "model_group" not in request_kwargs["metadata"] + redis_cache.async_get_cache.assert_awaited_once() + + @pytest.mark.asyncio + async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") + redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) + router._update_redis_cache(cache=redis_cache) + + main_response = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=self._request_kwargs(), + ) + subagent_response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert main_response is not None + assert main_response.model == "cheap-model" + assert subagent_response is not None + assert subagent_response.model == "cheap-model" + assert redis_cache.async_set_cache.await_count == 2 + @pytest.mark.asyncio async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): from types import SimpleNamespace From 1306a4505a03bdced3ccdf5ac6b26a719867bbbb Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:01:15 -0700 Subject: [PATCH 32/35] fix(router): preserve bound router fallbacks for subagents --- .../router_utils/fallback_event_handlers.py | 12 +++++----- .../test_fallback_event_handlers.py | 22 +++++++++++++++++++ tests/test_litellm/test_router.py | 21 ++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0167721f9fe..f7855cb38ff 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -231,8 +232,6 @@ def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selec on /v1/messages the top-level ``metadata`` dict is the provider's own request field, so a blanket write would forward the tier stamp upstream. """ - from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs - if request_kwargs is None: return bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) @@ -267,10 +266,13 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: """ Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, - and the requested group still resolves when no tier-keyed chain exists, so configs keyed - on the router name (the documented contract) keep working behind auto-routers. + then the routed group, then the requested group. The routed group differs when Claude Code + session affinity remaps a subagent's concrete model to its bound router. """ - ordered: Final = (get_pre_routing_selection(kwargs), model_group) + metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) + routed_group_value: Final = metadata.get("model_group") if isinstance(metadata, Mapping) else None + routed_group: Final = routed_group_value if isinstance(routed_group_value, str) else None + ordered: Final = (get_pre_routing_selection(kwargs), routed_group, model_group) return tuple(dict.fromkeys(group for group in ordered if group)) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 9e51a60364b..a4965c49f07 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1220,6 +1220,28 @@ class TestOrderedFallbackLookupGroups: assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) assert fallback_lookup_groups({}, None) == () + def test_session_remap_keeps_the_bound_router_between_tier_and_requested_group(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = { + "litellm_metadata": { + PRE_ROUTING_SELECTED_MODEL_KEY: "tier1", + "model_group": "smart-router", + } + } + + assert fallback_lookup_groups(kwargs, "requested-model") == ( + "tier1", + "smart-router", + "requested-model", + ) + assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ( + "requested-model", + ) + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): from litellm.router_utils.fallback_event_handlers import ( get_fallback_model_group_for_lookup_groups, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 16c4d997ad7..0a433d0b0a1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8679,6 +8679,27 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response.choices[0].message.content == "expensive response" assert subagent_kwargs["metadata"]["routing_decision"]["routed_model"] == "cheap-model" + @pytest.mark.asyncio + async def test_subagent_can_use_the_bound_router_name_fallback(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"smart-router": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **self._request_kwargs(agent_id="agent-1234"), + ) + + assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From f5721ff6f063cd6cfed4285a4bee80e6cf1fc81b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:17:35 -0700 Subject: [PATCH 33/35] test(router): cover Claude session multi-hop fallbacks --- tests/test_litellm/test_router.py | 59 +++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 0a433d0b0a1..db73f0da3c2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8700,6 +8700,65 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio + async def test_anthropic_subagent_four_fallback_hops_use_each_current_model_chain(self): + from litellm.types.router import TaggedPreRoutingStrategy + + failing_groups = ("cheap-model", "fallback-1", "fallback-2", "fallback-3") + router = litellm.Router( + model_list=[ + *( + { + "model_name": group, + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "litellm.RateLimitError", + }, + } + for group in failing_groups + ), + { + "model_name": "requested-model", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "requested response", + }, + }, + { + "model_name": "fallback-4", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "fourth fallback response", + }, + }, + ], + fallbacks=[ + {"smart-router": ["fallback-1"]}, + {"fallback-1": ["fallback-2"]}, + {"fallback-2": ["fallback-3"]}, + {"fallback-3": ["fallback-4"]}, + ], + num_retries=0, + max_fallbacks=4, + ) + router.complexity_routers = { + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=self._RewriteStrategy()),) + } + main_kwargs = self._request_kwargs() + main_kwargs["litellm_metadata"] = main_kwargs.pop("metadata") + await router.async_pre_routing_hook(model="smart-router", request_kwargs=main_kwargs) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + subagent_kwargs["litellm_metadata"] = subagent_kwargs.pop("metadata") + + response = await router.aanthropic_messages( + model="requested-model", + messages=[{"role": "user", "content": "subagent turn"}], + max_tokens=64, + **subagent_kwargs, + ) + + assert response["content"][0]["text"] == "fourth fallback response" + @pytest.mark.asyncio async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From 632898b2a662049f18d35c0f59ab1c0e348c8f44 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 16:51:43 -0700 Subject: [PATCH 34/35] feat(router): add a hybrid classifier that defers near tier boundaries (#39403) --- .../complexity_router/README.md | 43 ++++ .../complexity_router/complexity_router.py | 35 +++ .../complexity_router/config.py | 46 +++- litellm/types/utils.py | 1 + .../router_strategy/test_complexity_router.py | 201 ++++++++++++++++++ .../components/AutoRouters/autoRouterRows.ts | 1 + .../add_model/ClassificationMethodConfig.tsx | 47 ++++ .../add_model/ComplexityRouterConfig.tsx | 12 +- .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 37 +++- .../build_complexity_router_config.ts | 8 + .../src/components/add_model/tier_rows.ts | 4 +- ...d_updated_complexity_router_config.test.ts | 25 ++- .../edit_auto_router_modal.tsx | 5 + .../LogDetailsDrawer/RoutingDecisionCard.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 13 +- 16 files changed, 460 insertions(+), 20 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 8e0cad39561..ad8b67d5e8f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -274,6 +274,49 @@ except that the heuristic outcome is the one already computed rather than a seco Spend logs record `routing_decision.cause` as `heuristic_first_short_circuit` when the classifier was skipped, and `llm_classifier` when it ran, so the two are told apart per request. +### Hybrid + +`classifier_type: hybrid` also scores locally first, but it asks a different question than +`heuristic_first`. Where heuristic-first asks how CHEAP the scorer's tier is and pays for the +classifier on everything above a ceiling, hybrid asks how DECIDED the score is and pays for the +classifier only where the score lands near a tier boundary. A confident score keeps its tier at +every tier, the most expensive one included: + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: hybrid + hybrid_boundary_margin: 0.03 + classifier_llm_config: + model: gpt-4o-mini + tiers: + SIMPLE: gpt-4o-mini + MEDIUM: gpt-4o + COMPLEX: claude-sonnet-4 + REASONING: o1-preview +``` + +A request routes on the scorer's own tier when its score is further than `hybrid_boundary_margin` +from every active boundary. Everything else goes to the classifier: a score inside the band, where a +hair's difference would have named the adjacent tier and its model pool, and a prompt where no +dimension fired at all, which has no opinion to be confident about. `hybrid_boundary_margin` is +required for this type and rejected on the others, the same way `heuristic_first_max_tier` is +required for heuristic-first, so the two modes are told apart by the knob each one takes rather than +by a shared field that means something different per type. + +Pick the margin against the score distribution rather than by intuition. The scorer combines a small +set of discretely weighted dimensions, so achievable scores cluster on a lumpy grid instead of +spreading smoothly, and widening the margin admits whole clusters at once rather than a few more +requests. Spend logs record `routing_decision.cause` as `hybrid_short_circuit` when the classifier +was skipped and `llm_classifier` when it ran. + +Operator-defined tier sets (`tier_definitions`) are not supported here, for the same reason they are +not supported under heuristic-first: the scorer only produces the built-in tiers. Classifier failure +behaves exactly as it does under `classifier_type: llm`. + ### Reasoning Override If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a123a75dd81..d205db90607 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -799,6 +799,7 @@ class ClassificationOutcome(NamedTuple): "reasoning_override", "llm_classifier", "heuristic_first_short_circuit", + "hybrid_short_circuit", "housekeeping", "classifier_plugin", "classifier_fallback", @@ -1241,6 +1242,15 @@ class ComplexityRouter(CustomLogger): return tier, weighted_score, tuple(signals), "heuristic_scorer" + def _is_near_tier_boundary(self, score: float, margin: float) -> bool: + boundaries: Final = self._effective_tier_boundaries() + active_boundaries: Final = ( + boundaries["simple_medium"], + boundaries["medium_complex"], + boundaries["complex_reasoning"], + ) + return any(abs(score - boundary) <= margin for boundary in active_boundaries) + def _effective_reasoning_override_min_score(self) -> float: """The score a request must reach before the reasoning-marker override may promote it. @@ -1367,6 +1377,8 @@ class ComplexityRouter(CustomLogger): return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages) + if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None: + return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None: tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) @@ -1418,6 +1430,29 @@ class ComplexityRouter(CustomLogger): return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit") return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _classify_hybrid( + self, + prompt: str, + system_prompt: str | None, + request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is + messages: Sequence[Mapping[str, object]] | None, + ) -> ClassificationOutcome: + """Score locally, and only pay for the classifier when the score sits near a tier boundary. + + Where heuristic_first asks how CHEAP the scorer's tier is, this asks how DECIDED it is, so a + confident score keeps its tier at every tier including the most expensive one. Two things make + a score undecided: landing within hybrid_boundary_margin of an active boundary, where a + hair's difference in score would have named the adjacent tier and its model pool, and firing + no dimension at all, which scores 0.0 and lands SIMPLE by default rather than by evidence. + """ + tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) + scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) + margin: Final = self.config.hybrid_boundary_margin + decided: Final = margin is not None and bool(signals) and not self._is_near_tier_boundary(score, margin) + if decided: + return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="hybrid_short_circuit") + return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _llm_classifier_outcome( self, prompt: str, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fdf2a3a0b39..0ae0db63fad 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -43,7 +43,7 @@ DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubri # The classifier_type values that can call classifier_llm_config.model. Every consumer asking # "is the classifier model a real dependency of this router" resolves it here, including the ones # that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier. -LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first"}) +LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first", "hybrid"}) TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( @@ -627,12 +627,13 @@ class ComplexityRouterConfig(BaseModel): ) # Classifier strategy - classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first"] = Field( + classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first", "hybrid"] = Field( default="heuristic", description=( "Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, " - "an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays " - "for the LLM classifier when the local scorer does not confidently land a cheap tier" + "an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays " + "for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', " + "which trusts the local scorer everywhere except when its score lands near a tier boundary" ), ) heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field( @@ -644,7 +645,10 @@ class ComplexityRouterConfig(BaseModel): ) classifier_llm_config: ClassifierLLMConfig | None = Field( default=None, - description="Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first'", + description=( + "Configuration for the LLM classifier; required when classifier_type is 'llm', " + "'heuristic_first' or 'hybrid'" + ), ) heuristic_first_max_tier: str | None = Field( default=None, @@ -659,6 +663,19 @@ class ComplexityRouterConfig(BaseModel): "may not name the highest one, since that would make the LLM classifier unreachable." ), ) + hybrid_boundary_margin: float | None = Field( + default=None, + ge=0, + le=1, + description=( + "How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the " + "tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than " + "this from every active boundary routes on the scorer's own tier with no classifier call, at any " + "tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A " + "prompt where no dimension fired still goes to the classifier, since the scorer has no opinion " + "to be near a boundary with. 0 escalates only scores sitting exactly on a boundary." + ), + ) classifier_plugin: ClassifierPlugin | None = Field( default=None, description=( @@ -1135,6 +1152,23 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_hybrid_boundary_margin(self) -> "ComplexityRouterConfig": + if self.classifier_type != "hybrid": + if self.hybrid_boundary_margin is not None: + raise ValueError( + f"hybrid_boundary_margin is set but classifier_type is {self.classifier_type!r}; " + "the scorer would never consult the classifier on a near-boundary score. Set " + "classifier_type 'hybrid' or remove hybrid_boundary_margin" + ) + return self + if self.hybrid_boundary_margin is None: + raise ValueError( + "hybrid_boundary_margin is required when classifier_type is 'hybrid': without a margin no " + "score is ever near enough to a boundary to escalate, which is classifier_type 'heuristic'" + ) + return self + @field_validator("fallback_tier") @classmethod def _reject_blank_optional_text(cls, value: str | None) -> str | None: @@ -1257,7 +1291,7 @@ class ComplexityRouterConfig(BaseModel): ) if duplicated: raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}") - if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first"): + if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"): raise ValueError( "tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only " "produces the four built-in tiers, as does heuristic_v2" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 09c01873a9d..ee6f09e05dc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2852,6 +2852,7 @@ RoutingDecisionCause = Literal[ # scorer, and from "classifier_fallback", which is the scorer running because a call failed: # only this cause means an LLM classifier was configured, reachable, and deliberately skipped. "heuristic_first_short_circuit", + "hybrid_short_circuit", # The operator's classifier plugin (classifier_type 'custom') decided the tier. "classifier_plugin", # The LLM classifier or classifier plugin failed on a router with an operator-defined diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 371631d6297..941d78085e4 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10089,6 +10089,207 @@ class TestHeuristicFirst: assert outcome.cause == "default_model_fallback" +# Scores 0.175 with one signal, so it sits 0.025 from simple_medium: the pair of tiers either side of +# that boundary are different model pools, and a hair's difference in score picks the other one. +NEAR_BOUNDARY_PROMPT = ( + "design a distributed cache with consistent hashing, then explain the failure modes step by step" +) + +# Scores 0.075 with signals, the far side of any margin under 0.075: the scorer is decided here. +CLEAR_OF_BOUNDARY_PROMPT = "explain step by step how consistent hashing rebalances keys" + + +def _hybrid_router(mock_router_instance, **config_overrides): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "tier_boundaries": dict(HEURISTIC_FIRST_BOUNDARIES), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400}, + **config_overrides, + } + return ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + +class TestHybridConfig: + """Config validation for classifier_type='hybrid'.""" + + @pytest.mark.parametrize( + "overrides, expected", + [ + ({"classifier_llm_config": None}, "classifier_llm_config is required"), + ({"hybrid_boundary_margin": None}, "hybrid_boundary_margin is required"), + ({"hybrid_boundary_margin": -0.01}, "greater than or equal to 0"), + ({"hybrid_boundary_margin": 1.01}, "less than or equal to 1"), + ], + ) + def test_rejects_incoherent_config(self, overrides, expected): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier"}, + **overrides, + } + with pytest.raises(ValidationError, match=expected): + ComplexityRouterConfig(**config) + + @pytest.mark.parametrize("classifier_type", ["heuristic", "llm", "custom", "heuristic_first"]) + def test_margin_rejected_on_every_other_classifier_type(self, classifier_type): + """A margin on a router that never compares a score to a boundary is a silent no-op, so it is + refused rather than accepted and ignored. heuristic_first is in this list on purpose: its + ceiling is a different question from proximity, and accepting both on one router would make + two modes out of one classifier_type.""" + config: dict[str, object] = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": classifier_type, + "hybrid_boundary_margin": 0.03, + } + if classifier_type in ("llm", "heuristic_first"): + config["classifier_llm_config"] = {"model": "haiku-classifier"} + if classifier_type == "heuristic_first": + config["heuristic_first_max_tier"] = "SIMPLE" + if classifier_type == "custom": + config["classifier_plugin"] = _FixedTierClassifier("SIMPLE") + with pytest.raises(ValidationError, match="hybrid_boundary_margin is set but classifier_type"): + ComplexityRouterConfig(**config) + + def test_the_cheap_tier_ceiling_is_rejected_here(self): + """The two modes are told apart by which knob they take, so the ceiling is refused on hybrid + exactly as the margin is refused on heuristic_first.""" + with pytest.raises(ValidationError, match="heuristic_first_max_tier is set but classifier_type"): + ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + heuristic_first_max_tier="SIMPLE", + classifier_llm_config={"model": "haiku-classifier"}, + ) + + def test_custom_tier_set_is_rejected(self): + """The scorer only emits the four built-in tiers, so it cannot judge proximity on a replaced set.""" + with pytest.raises(ValidationError, match="tier_definitions requires classifier_type"): + ComplexityRouterConfig( + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + tier_definitions=[{"name": "lo", "description": "x"}, {"name": "hi", "description": "y"}], + tiers={"lo": "gpt-4o-mini", "hi": "gpt-4o"}, + ) + + def test_classifier_model_is_a_dependency(self): + config = ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + ) + assert config.uses_llm_classifier is True + + +class TestHybrid: + """Behavior of the hybrid chain: the scorer keeps its tier unless the score is near a boundary.""" + + @pytest.mark.asyncio + async def test_near_boundary_prompt_escalates(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + _tier, score, signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + assert signals and abs(score - HEURISTIC_FIRST_BOUNDARIES["simple_medium"]) < 0.03 + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.tier == ComplexityTier.COMPLEX + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_score_clear_of_every_boundary_keeps_the_heuristic_tier(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router(mock_router_instance) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.SIMPLE + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_an_expensive_tier_short_circuits_too(self, mock_router_instance): + """This is the whole difference from heuristic_first, which would have escalated this by tier + alone. Hybrid asks whether the score is DECIDED, not whether the tier is cheap.""" + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": -0.9, "medium_complex": -0.8, "complex_reasoning": -0.7}, + ) + + tier, _score, signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert (tier, bool(signals)) == (ComplexityTier.REASONING, True) + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.REASONING + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_widening_the_margin_escalates_what_a_narrow_one_kept(self, mock_router_instance): + """The margin is the knob: the same prompt short-circuits at 0.03 and escalates at 0.08.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + router = _hybrid_router(mock_router_instance, hybrid_boundary_margin=0.08) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_a_zero_margin_escalates_only_an_exact_boundary_score(self, mock_router_instance): + """0 is a real margin, not an off switch: a score sitting exactly on the line still escalates. + + The boundary is spelled as the scorer's own accumulated float rather than the 0.075 it prints + as, because the comparison is on raw floats: a boundary written 0.075 sits 1.4e-17 away from + this score and a zero margin correctly declines to call that exact.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + on_the_line = 0.07499999999999998 + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": on_the_line, "medium_complex": 0.35, "complex_reasoning": 0.60}, + hybrid_boundary_margin=0, + ) + + _tier, score, _signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert score == on_the_line + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_no_signal_prompt_escalates_however_far_from_a_boundary(self, mock_router_instance): + """The scorer with no opinion has no tier to be confident about, so proximity cannot save it.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + tier, score, signals, _cause = router._score_and_classify(NO_SIGNAL_PROMPT) + assert (tier, score, signals) == (ComplexityTier.SIMPLE, 0.0, ()) + + outcome = await router.aclassify(NO_SIGNAL_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_classifier_failure_falls_back_to_the_scorer(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier exploded")) + router = _hybrid_router(mock_router_instance) + expected_tier, expected_score, expected_signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + + assert (outcome.tier, outcome.score, outcome.signals) == (expected_tier, expected_score, expected_signals) + assert outcome.cause == "heuristic_scorer" + + def _windowed_router(*deployments: tuple) -> Router: """Real Router; each deployment is (group, provider_model, declared window or None). None means no declared override on a model the cost map does not know: unresolvable.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index bbdf4697315..c4d7f45b7cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -58,6 +58,7 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models)); const COMPLEXITY_TYPE_LABELS: Record = { llm: "LLM Classifier", heuristic_first: "Heuristic first", + hybrid: "Hybrid", custom: "Custom classifier", }; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index c2ed4b1f55a..a29527a20fa 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -35,6 +35,7 @@ import { heuristicScoringRole, usesLlmClassifier, DEFAULT_HEURISTIC_FIRST_MAX_TIER, + DEFAULT_HYBRID_BOUNDARY_MARGIN, HEURISTIC_FIRST_MAX_TIER_KEYS, effectiveClassifierType, } from "./ComplexityRouterConfig"; @@ -50,6 +51,7 @@ const HEURISTIC_V2_EXPLANATION = const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; +const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -213,6 +215,18 @@ const ClassifierTypeRadios: React.FC<{ + + + ); @@ -263,6 +277,8 @@ const ClassificationMethodConfig: React.FC = ({ classifierType === "heuristic_first" ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER : undefined, + hybrid_boundary_margin: + classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined, }; onChange(nextValue); }; @@ -271,6 +287,13 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, heuristic_first_max_tier: tier }); }; + const handleHybridBoundaryMarginChange = (raw: string) => { + setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ ...value, hybrid_boundary_margin: Math.min(1, Math.max(0, parsed)) }); + }; + const handleClassificationPromptChange = (classificationPrompt: string | undefined) => { onChange({ ...value, classification_prompt: classificationPrompt }); }; @@ -391,6 +414,30 @@ const ClassificationMethodConfig: React.FC = ({ )} + {classifierType === "hybrid" && ( +
+ Boundary margin + handleHybridBoundaryMarginChange(event.target.value)} + onBlur={() => setDraft(null)} + className="w-full" + /> +

+ A score further than this from every tier boundary routes on the scorer's own tier, however expensive + that tier is. A score closer than this, and anything the scorer found no signal for at all, goes to the + classifier to break the tie +

+
+ )} +
How often to classify - classifierType === "llm" || classifierType === "heuristic_first"; + classifierType === "llm" || classifierType === "heuristic_first" || classifierType === "hybrid"; export type ClassifierFallback = "heuristic" | "default_model"; @@ -162,7 +162,8 @@ export const heuristicScoringRoleFor = ( classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { if (classifierType === "heuristic_v2") return "never"; - if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides"; + if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") + return "decides"; return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never"; }; @@ -404,6 +405,8 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */ + hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; session_affinity?: boolean; modality_routing?: boolean; @@ -516,6 +519,9 @@ export const effectiveTierLabel = (tier: keyof ComplexityTiers, tierLabels: Comp export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; +/** What the Hybrid radio starts at. Required by that type, so the form always has a value to send. */ +export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03; + /** * Tiers the heuristic_first threshold may name. The top tier is excluded because it would short * circuit every request and leave the classifier unreachable, which the backend rejects. diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 03b7432e4ec..81acb7b7e50 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 2d8f55b2b64..c6370541a45 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -817,6 +817,39 @@ describe("heuristic_first", () => { }); }); +describe("hybrid", () => { + const hybridParams: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "hybrid", + hybridBoundaryMargin: 0.03, + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierFallback: "default_model", + }; + + it("emits hybrid_boundary_margin, zero included since exactly-on-a-boundary is a real setting", () => { + expect(buildComplexityRouterConfig(hybridParams).hybrid_boundary_margin).toBe(0.03); + expect(buildComplexityRouterConfig({ ...hybridParams, hybridBoundaryMargin: 0 }).hybrid_boundary_margin).toBe(0); + }); + + it("keeps every classifier key the operator set, since hybrid still calls the classifier", () => { + const config = buildComplexityRouterConfig(hybridParams); + expect(config.classifier_type).toBe("hybrid"); + expect(config.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); + expect(config.classifier_fallback).toBe("default_model"); + }); + + it("omits hybrid_boundary_margin on every other classifier type, which the backend rejects it on", () => { + for (const classifierType of ["heuristic", "llm", "heuristic_first"] as const) { + const config = buildComplexityRouterConfig({ + ...hybridParams, + classifierType, + ...(classifierType === "heuristic_first" && { heuristicFirstMaxTier: "SIMPLE" }), + }); + expect(config.hybrid_boundary_margin).toBeUndefined(); + } + }); +}); + describe("classification_mode", () => { it("emits user_turn", () => { const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); @@ -924,10 +957,12 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { dimensionWeights: { length: 1 }, reasoningOverrideMinScore: 0.5, heuristicFirstMaxTier: "SIMPLE", + hybridBoundaryMargin: 0.03, customTechnicalKeywords: ["kubernetes"], }; const emittingType = key === "heuristic_first_max_tier" ? "heuristic_first" : "llm"; - expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: emittingType })).toHaveProperty(key); + const typeForKey = key === "hybrid_boundary_margin" ? "hybrid" : emittingType; + expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: typeForKey })).toHaveProperty(key); expect(build(loaded)).not.toHaveProperty(key); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index af34dc92c0f..94c5badf6d1 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -107,6 +107,7 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + hybridBoundaryMargin?: number; classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; modalityRouting?: boolean; @@ -163,6 +164,7 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; @@ -352,6 +354,7 @@ const classifierWireFields = ( classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, @@ -360,6 +363,7 @@ const classifierWireFields = ( | "classifierLlmConfig" | "classifierFallback" | "heuristicFirstMaxTier" + | "hybridBoundaryMargin" | "classifierContextWindowSize" | "classifierContextBudgetChars" | "classifierContextIncludeAssistantTurns" @@ -371,6 +375,8 @@ const classifierWireFields = ( classifierFallback !== undefined && { classifier_fallback: classifierFallback }), ...(effectiveType === "heuristic_first" && heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }), + ...(effectiveType === "hybrid" && + hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }), ...(usesLlmClassifier(effectiveType) && classifierContextWindowSize !== undefined && { classifier_context_window_size: classifierContextWindowSize, @@ -399,6 +405,7 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + hybridBoundaryMargin, classificationMode, sessionAffinity, modalityRouting, @@ -444,6 +451,7 @@ export const buildComplexityRouterConfig = ({ classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, diff --git a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts index d2706a82037..3c5b149f4da 100644 --- a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts +++ b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts @@ -122,10 +122,10 @@ export const CUSTOM_TIER_RESTRICTIONS = { reason: "Session pinning escalates along the built-in tier ladder, which your tier set replaces", }, heuristicClassifier: { - omit: ["heuristic_first_max_tier"], + omit: ["heuristic_first_max_tier", "hybrid_boundary_margin"], reason: "The heuristic scorer only produces the built-in tiers, so an edited set needs the LLM classifier. " + - "Heuristic first is out for the same reason: its local scorer decides the cheap traffic", + "Heuristic first and hybrid are out for the same reason: their local scorer decides the traffic it is sure of", }, heuristicScoring: { omit: [ diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 3e63cbd3b32..fae6744d3ac 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -520,16 +520,21 @@ describe("managed keys survive an untouched open-and-save", () => { }; // tier_definitions, fallback_tier and classification_prompt cannot sit beside heuristic_first, which - // this fixture uses, so no single stored config can hold every managed key. They get their own round - // trip below. - const CUSTOM_TIER_ONLY_KEYS = new Set(["tier_definitions", "fallback_tier", "classification_prompt"]); + // this fixture uses, and hybrid_boundary_margin belongs to the sibling hybrid type, so no single + // stored config can hold every managed key. Each gets its own round trip below. + const KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS = new Set([ + "tier_definitions", + "fallback_tier", + "classification_prompt", + "hybrid_boundary_margin", + ]); it("carries every managed key a built-in router can hold through hydrate then save", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); const saved = buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated); const dropped = [...MANAGED_COMPLEXITY_ROUTER_KEYS] - .filter((key) => !CUSTOM_TIER_ONLY_KEYS.has(key)) + .filter((key) => !KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS.has(key)) .filter((key) => saved[key] === undefined); expect(dropped).toEqual([]); }); @@ -583,6 +588,18 @@ describe("managed keys survive an untouched open-and-save", () => { ); }); + it("round-trips a hybrid router's margin, which save requires and the backend rejects without", () => { + const storedHybrid: Record = { + ...STORED_ALL_MANAGED, + classifier_type: "hybrid", + hybrid_boundary_margin: 0.05, + }; + delete storedHybrid.heuristic_first_max_tier; + const hydrated = hydrateComplexityRouterConfig(storedHybrid, undefined); + expect(hydrated.hybrid_boundary_margin).toBe(0.05); + expect(buildUpdatedComplexityRouterConfig(storedHybrid, hydrated).hybrid_boundary_margin).toBe(0.05); + }); + it("round-trips the heuristic_first threshold, which save requires and the backend rejects without", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); expect(hydrated.heuristic_first_max_tier).toBe("SIMPLE"); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e18582f77a0..7aa5f4beb1f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -89,6 +89,7 @@ export interface StoredComplexityRouterConfig { plan_mode_min_tier?: unknown; classification_prompt?: unknown; heuristic_first_max_tier?: unknown; + hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; @@ -167,6 +168,8 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, classification_mode: parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" ? parsedConfig.classification_mode @@ -214,6 +217,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "hybrid_boundary_margin", "classification_mode", "session_affinity", "modality_routing", @@ -303,6 +307,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index 66e80dcc655..ffef01010b7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -86,6 +86,7 @@ const CONSTANT_CAUSE_LABELS: Record = { heuristic_scorer: "Heuristic scorer", heuristic_v2: "Heuristic v2", heuristic_first_short_circuit: "Heuristic scorer, classifier skipped", + hybrid_short_circuit: "Heuristic scorer, score clear of every boundary", classifier_plugin: "Custom classifier plugin", semantic_keyword_match: "Semantic keyword match", session_affinity_pin: "Pinned to session", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79e44de4f1f..be81450b941 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34560,7 +34560,7 @@ export interface components { * @enum {string} */ classifier_fallback: "heuristic" | "default_model"; - /** @description Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first' */ + /** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'heuristic_first' or 'hybrid' */ classifier_llm_config?: components["schemas"]["ClassifierLLMConfig"] | null; /** * Classifier Plugin @@ -34575,11 +34575,11 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier + * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary * @default heuristic * @enum {string} */ - classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first"; + classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first" | "hybrid"; /** * Code Keywords * @description Keywords indicating code-related content @@ -34651,6 +34651,11 @@ export interface components { * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; + /** + * Hybrid Boundary Margin + * @description How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than this from every active boundary routes on the scorer's own tier with no classifier call, at any tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A prompt where no dimension fired still goes to the classifier, since the scorer has no opinion to be near a boundary with. 0 escalates only scores sitting exactly on a boundary. + */ + hybrid_boundary_margin?: number | null; /** * Keyword Tier Rules * @description Rules that force a specific tier when their keywords match the prompt @@ -35867,7 +35872,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ From b0fe71010b89425501587a3e7e66ede9a4e6e775 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:56:35 -0700 Subject: [PATCH 35/35] fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk (#39010) * fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama_chat): trim finish_reason override comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mubashir1osmani --- litellm/llms/ollama/chat/transformation.py | 7 +- .../test_streaming_iterator_stop_reason.py | 79 +++++++++++++++++++ .../ollama/test_ollama_chat_transformation.py | 40 ++++++++++ 3 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index de626b468f0..181894646e3 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -423,6 +423,7 @@ class OllamaChatConfig(BaseConfig): class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): started_reasoning_content: bool = False finished_reasoning_content: bool = False + seen_tool_calls: bool = False def _is_function_call_complete(self, function_args: str | dict) -> bool: if isinstance(function_args, dict): @@ -468,6 +469,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): # process tool calls - if complete function arg - add id to tool call tool_calls: Final = chunk["message"].get("tool_calls") if tool_calls is not None: + self.seen_tool_calls = True for tool_call in tool_calls: function_args = tool_call.get("function").get("arguments") if function_args is not None and len(function_args) > 0: @@ -508,9 +510,10 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): if chunk["done"] is True: finish_reason = chunk.get("done_reason") or "stop" - # Override finish_reason when tool_calls are present + # Override finish_reason when tool_calls appeared in any chunk # Fixes: https://github.com/BerriAI/litellm/issues/18922 - if tool_calls is not None: + # Fixes: https://github.com/BerriAI/litellm/issues/34692 + if self.seen_tool_calls: finish_reason = "tool_calls" choices = [ StreamingChoices( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py new file mode 100644 index 00000000000..4b95b36fec3 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py @@ -0,0 +1,79 @@ +""" +Regression tests for issue #34692. + +ollama_chat streams tool_calls in a mid-stream chunk while its final +(``done: true``) chunk carries only ``done_reason: "stop"``. The provider +iterator must remember the earlier tool_calls and stamp +``finish_reason="tool_calls"`` on the final chunk, so the Anthropic +``/v1/messages`` bridge emits ``stop_reason: "tool_use"``. Before the fix the +bridge emitted ``stop_reason: "end_turn"`` and Anthropic tool-runners +(Claude Code, ``messages.stream``) silently dropped the tool call. +""" + +import pytest + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.llms.ollama.chat.transformation import ( + OllamaChatCompletionResponseIterator, +) +from litellm.types.utils import ModelResponseStream + +_OLLAMA_TOOL_CHUNK = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}}], + }, + "done": False, +} +_OLLAMA_DONE_CHUNK = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 100, + "eval_count": 20, +} + + +def _ollama_streamed_chunks() -> list[ModelResponseStream]: + iterator = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + return [iterator.chunk_parser(_OLLAMA_TOOL_CHUNK), iterator.chunk_parser(_OLLAMA_DONE_CHUNK)] + + +class _AsyncStream: + def __init__(self, items: list[ModelResponseStream]): + self._it = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration: + raise StopAsyncIteration + + +def _assert_tool_use_stop_reason(events: list[dict]) -> None: + block_types = [e["content_block"]["type"] for e in events if e.get("type") == "content_block_start"] + assert "tool_use" in block_types, f"no tool_use content block opened: {events}" + message_deltas = [e for e in events if e.get("type") == "message_delta"] + assert message_deltas, f"no message_delta emitted: {events}" + assert message_deltas[-1]["delta"]["stop_reason"] == "tool_use", ( + f"expected stop_reason 'tool_use', got: {message_deltas[-1]}" + ) + + +def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_sync(): + wrapper = AnthropicStreamWrapper(completion_stream=iter(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason(list(wrapper)) + + +@pytest.mark.asyncio +async def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_async(): + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason([event async for event in wrapper]) diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index 8f3dbf7b0d9..25f9645faa0 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -615,6 +615,46 @@ class TestOllamaFinishReasonLength: result.choices[0].finish_reason == "stop" ), f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" + def test_finish_reason_tool_calls_streamed_before_done_chunk(self): + """Streaming: tool_calls arriving mid-stream (not on the done chunk) must + still produce finish_reason='tool_calls' on the final chunk. + + Regression test for https://github.com/BerriAI/litellm/issues/34692: + Ollama emits tool_calls in an earlier chunk and the done chunk carries + none, which left finish_reason at 'stop' and made the Anthropic + /v1/messages bridge emit stop_reason 'end_turn' instead of 'tool_use'. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + tool_chunk = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + {"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}} + ], + }, + "done": False, + } + done_chunk = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + } + + tool_result = iterator.chunk_parser(tool_chunk) + assert tool_result.choices[0].delta.tool_calls is not None + + done_result = iterator.chunk_parser(done_chunk) + assert ( + done_result.choices[0].finish_reason == "tool_calls" + ), f"Expected 'tool_calls' when tool_calls were streamed earlier, got '{done_result.choices[0].finish_reason}'" + class TestOllamaReasoningContentStreaming: """Test that reasoning_content is properly extracted from all thinking chunks."""