From b61484e6c919b2d8718249ef10889029192fa5b5 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:31:50 +0000 Subject: [PATCH 01/88] feat: add Amazon Bedrock AgentCore Web Search as a native search provider Adds 'agentcore' to SearchProviders, backed by an AgentCore Gateway web-search connector target (MCP tools/call over Streamable HTTP). Web Search on Amazon Bedrock AgentCore is an AWS-managed web index (GA June 2026). Exposing it as a native search provider lets Bedrock users enable Claude Code / Anthropic-native WebSearch through websearch_interception with a pure-YAML config and AWS-native auth, keeping the whole search path inside AWS. Implementation: - New AgentCoreSearchConfig (litellm/llms/bedrock/search/) reusing BaseAWSLLM credential resolution. Auth follows the gateway's inbound authorizer type: AWS_IAM gateways get a SigV4-signed request (explicit aws_access_key_id/aws_secret_access_key params or the default credential chain); CUSTOM_JWT gateways get an OAuth2 bearer token via api_key / AGENTCORE_GATEWAY_TOKEN - SigV4 signing region is derived from the gateway URL so callers don't need aws_region_name to match their default region - Adds an optional sign_request() hook to BaseSearchConfig (no-op by default) and teaches the search HTTP handler to send a signed body verbatim, mirroring the existing anthropic_messages/chat pattern - Handles both plain-JSON and SSE-framed MCP responses, propagates MCP errors, truncates queries to the 200-char gateway limit Tested: - 13 unit tests: payload/signing, explicit AKSK passthrough, bearer token via api_key and env, query truncation, SSE frames, MCP error propagation, region derivation - Verified end-to-end against real AWS_IAM and CUSTOM_JWT gateways, including full Claude Code CLI WebSearch round-trips through the proxy with websearch_interception --- .../llms/base_llm/search/transformation.py | 23 ++ litellm/llms/bedrock/search/__init__.py | 0 litellm/llms/bedrock/search/transformation.py | 256 ++++++++++++++++++ litellm/llms/custom_httpx/llm_http_handler.py | 34 +++ .../agentcore_websearch_config.yaml | 39 +++ litellm/types/utils.py | 1 + litellm/utils.py | 2 + tests/search_tests/test_agentcore_search.py | 231 ++++++++++++++++ 8 files changed, 586 insertions(+) create mode 100644 litellm/llms/bedrock/search/__init__.py create mode 100644 litellm/llms/bedrock/search/transformation.py create mode 100644 litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml create mode 100644 tests/search_tests/test_agentcore_search.py diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index fdfac6f5f9f..7a93cf43ca7 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -178,6 +178,29 @@ class BaseSearchConfig: """ return headers + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: Union[dict, list[dict]], + api_base: str, + api_key: str | None = None, + ) -> tuple[dict, bytes | None]: + """ + OPTIONAL + + Sign the request. Providers like Bedrock AgentCore need to SigV4-sign + the request before sending it to the API. + + For all other providers, this is a no-op and we just return the headers. + + Returns: + Tuple of (headers, signed_json_body). When signed_json_body is not + None, the handler MUST send it verbatim as the request body — + re-serializing the payload would invalidate the signature. + """ + return headers, None + def get_complete_url( self, api_base: Optional[str], diff --git a/litellm/llms/bedrock/search/__init__.py b/litellm/llms/bedrock/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py new file mode 100644 index 00000000000..16d671f26b0 --- /dev/null +++ b/litellm/llms/bedrock/search/transformation.py @@ -0,0 +1,256 @@ +""" +Calls an Amazon Bedrock AgentCore Gateway web-search target (MCP protocol) to search the web. + +Web Search on Amazon Bedrock AgentCore exposes Amazon's managed web index through +an AgentCore Gateway MCP endpoint. + +AWS docs: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-web-search-tool.html + +Authentication (matches the gateway's inbound authorizer type): +- AWS_IAM gateway: the request is SigV4-signed. Credentials come from explicit + params (aws_access_key_id / aws_secret_access_key / aws_session_token / + aws_region_name — also settable in a proxy search_tools entry) or the + standard AWS credential chain (env / profile / IRSA / assumed role) +- CUSTOM_JWT gateway: pass the OAuth2 bearer token (e.g. Cognito + client_credentials) as api_key, or set AGENTCORE_GATEWAY_TOKEN + +Setup: + 1. Create an AgentCore Gateway with a web-search connector target + 2. Set AGENTCORE_GATEWAY_URL (or pass api_base) to the gateway MCP endpoint, e.g. + https://.gateway.bedrock-agentcore..amazonaws.com/mcp + 3. AWS_IAM: ensure the credentials allow bedrock-agentcore:InvokeGateway + CUSTOM_JWT: set AGENTCORE_GATEWAY_TOKEN (or pass api_key) + +Usage: + response = litellm.search( + query="latest AI developments", + search_provider="agentcore", + max_results=5, + aws_access_key_id="...", # optional — omit to use the default chain + aws_secret_access_key="...", + ) +""" + +import json +import re +from typing import Union + +import httpx + +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.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.secret_managers.main import get_secret_str + +# AgentCore web-search rejects queries longer than 200 characters +AGENTCORE_MAX_QUERY_LENGTH = 200 + +# Default MCP tool name for a gateway web-search connector target: +# "___". Override with AGENTCORE_SEARCH_TOOL_NAME +# or optional_params["tool_name"] when the target uses a custom name. +AGENTCORE_DEFAULT_TOOL_NAME = "web-search-tool___WebSearch" + + +class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): + def __init__(self) -> None: + BaseSearchConfig.__init__(self) + BaseAWSLLM.__init__(self) + + @staticmethod + def ui_friendly_name() -> str: + return "Web Search on Amazon Bedrock" + + def validate_environment( + self, + headers: dict, + api_key: str | None = None, + api_base: str | None = None, + **kwargs, + ) -> dict: + """ + Set MCP transport headers. Per the MCP Streamable HTTP transport spec, + the client MUST accept both application/json and text/event-stream. + + Authentication itself happens in sign_request(): bearer token for + CUSTOM_JWT gateways, AWS SigV4 for AWS_IAM gateways. + """ + headers["Content-Type"] = "application/json" + headers["Accept"] = "application/json, text/event-stream" + return headers + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict, + data: Union[dict, list[dict]] | None = None, + **kwargs, + ) -> str: + api_base = api_base or get_secret_str("AGENTCORE_GATEWAY_URL") + if not api_base: + raise ValueError( + "AGENTCORE_GATEWAY_URL is not set. Set it to your AgentCore Gateway MCP " + "endpoint (https://.gateway.bedrock-agentcore." + ".amazonaws.com/mcp) or pass api_base." + ) + return api_base + + def transform_search_request( + self, + query: Union[str, list[str]], + optional_params: dict, + **kwargs, + ) -> dict: + """ + Transform Search request to an MCP tools/call request. + + Args: + query: Search query (string or list of strings). AgentCore only + supports single string queries; lists are joined with spaces. + optional_params: Optional parameters for the request + - max_results: Maximum number of results (1-25), default 10 + - tool_name: Override the MCP tool name of the gateway target + + Returns: + Dict with the JSON-RPC 2.0 request body + """ + if isinstance(query, list): + query = " ".join(query) + query = query[:AGENTCORE_MAX_QUERY_LENGTH] + + tool_name = ( + optional_params.get("tool_name") + or get_secret_str("AGENTCORE_SEARCH_TOOL_NAME") + or AGENTCORE_DEFAULT_TOOL_NAME + ) + + arguments: dict[str, Union[str, int]] = {"query": query} + if "max_results" in optional_params: + arguments["maxResults"] = optional_params["max_results"] + + return { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": tool_name, "arguments": arguments}, + } + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: Union[dict, list[dict]], + api_base: str, + api_key: str | None = None, + ) -> tuple[dict, bytes | None]: + """ + Authenticate the MCP request. + + CUSTOM_JWT gateways: attach the caller's OAuth2 bearer token (api_key + or AGENTCORE_GATEWAY_TOKEN) — no AWS credentials involved. + + AWS_IAM gateways: SigV4-sign with the bedrock-agentcore service name. + """ + if not isinstance(request_data, dict): + raise ValueError("AgentCore search expects a single dict request body") + + bearer_token = api_key or get_secret_str("AGENTCORE_GATEWAY_TOKEN") + if bearer_token: + headers["Authorization"] = f"Bearer {bearer_token}" + return headers, json.dumps(request_data).encode() + + # The signing region must match the gateway's region — derive it from + # the gateway URL so callers don't have to set aws_region_name to a + # region different from their default. + signing_params = dict(optional_params) + if signing_params.get("aws_region_name") is None: + match = re.search( + r"\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com", + api_base, + ) + if match: + signing_params["aws_region_name"] = match.group(1) + + return self._sign_request( + service_name="bedrock-agentcore", + headers=headers, + optional_params=signing_params, + request_data=request_data, + api_base=api_base, + ) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> SearchResponse: + """ + Transform an MCP tools/call response to LiteLLM unified SearchResponse. + + The gateway returns JSON-RPC (as plain JSON or a single-message SSE + stream) whose result.content[] text blocks contain a JSON list of + {title, url, date/publishedDate, text} entries. + """ + response_json = self._parse_mcp_body(raw_response) + + if "error" in response_json: + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=f"AgentCore gateway MCP error: {response_json['error']}", + ) + + results: list[SearchResult] = [] + for block in response_json.get("result", {}).get("content", []): + if block.get("type") != "text": + continue + try: + parsed = json.loads(block["text"]) + except (json.JSONDecodeError, TypeError): + continue + items = parsed.get("results", []) if isinstance(parsed, dict) else parsed + for item in items: + if not isinstance(item, dict): + continue + results.append( + SearchResult( + title=item.get("title") or "", + url=item.get("url") or "", + snippet=item.get("text") or item.get("snippet") or "", + date=item.get("publishedDate") or item.get("date"), + last_updated=None, + ) + ) + + return SearchResponse(results=results, object="search") + + @staticmethod + def _parse_mcp_body(raw_response: httpx.Response) -> dict: + """Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response.""" + text = raw_response.text + if text.lstrip().startswith(("event:", "data:")): + for line in text.splitlines(): + if line.startswith("data:"): + return json.loads(line[len("data:") :].strip()) + raise BedrockError( + status_code=502, + message=f"AgentCore gateway returned SSE without a data frame: {text[:200]}", + ) + return raw_response.json() + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict, + ) -> Exception: + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3e6f9ee08ee..aa33865df43 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1737,6 +1737,15 @@ class BaseLLMHTTPHandler: api_key=api_key, ) + # Sign the request if the provider requires it (e.g. AWS SigV4) + headers, signed_json_body = provider_config.sign_request( + headers=headers, + optional_params=optional_params, + request_data=data, + api_base=complete_url, + api_key=api_key, + ) + ## LOGGING logging_obj.pre_call( input=query if isinstance(query, str) else str(query), @@ -1762,6 +1771,14 @@ class BaseLLMHTTPHandler: url=complete_url, headers=headers, ) + elif signed_json_body is not None: + # Send the signed body verbatim — re-serializing would break the signature + response = client.post( + url=complete_url, + headers=headers, + data=signed_json_body, + timeout=timeout, + ) else: # Make POST request with JSON data response = client.post( @@ -1821,6 +1838,15 @@ class BaseLLMHTTPHandler: api_key=api_key, ) + # Sign the request if the provider requires it (e.g. AWS SigV4) + headers, signed_json_body = provider_config.sign_request( + headers=headers, + optional_params=optional_params, + request_data=data, + api_base=complete_url, + api_key=api_key, + ) + ## LOGGING logging_obj.pre_call( input=query if isinstance(query, str) else str(query), @@ -1851,6 +1877,14 @@ class BaseLLMHTTPHandler: url=complete_url, headers=headers, ) + elif signed_json_body is not None: + # Send the signed body verbatim — re-serializing would break the signature + response = await async_httpx_client.post( + url=complete_url, + headers=headers, + data=signed_json_body, + timeout=timeout, + ) else: # Make async POST request with JSON data response = await async_httpx_client.post( diff --git a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml new file mode 100644 index 00000000000..f2c5a460bf0 --- /dev/null +++ b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml @@ -0,0 +1,39 @@ +# Claude Code / Anthropic-native web search on Bedrock, backed by +# Amazon Bedrock AgentCore Web Search (AWS-managed web index, no third-party +# search API). See litellm/llms/bedrock/search/transformation.py for details. + +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 + aws_region_name: us-east-1 + +search_tools: + - search_tool_name: agentcore-search + litellm_params: + search_provider: agentcore + # Your AgentCore Gateway MCP endpoint (gateway must have a `web-search` + # connector target). Alternatively set the AGENTCORE_GATEWAY_URL env var. + api_base: https://.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp + + # The gateway exposes the connector as "___WebSearch". + # Default is "web-search-tool___WebSearch", matching the target name used + # in the AWS docs' boto3/CLI setup examples. Set this ONLY if your target + # was created with a different name (misconfiguration surfaces as an MCP + # "tool not found" error): + # tool_name: MyWebSearchTarget___WebSearch + + # AWS_IAM gateway (default): SigV4-signed. Omit keys to use the standard + # AWS credential chain (env / profile / IRSA / instance role), or set them + # explicitly: + # aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + # aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + + # CUSTOM_JWT gateway alternative — OAuth2 bearer token instead of SigV4: + # api_key: os.environ/AGENTCORE_GATEWAY_TOKEN + +litellm_settings: + callbacks: ["websearch_interception"] + websearch_interception_params: + enabled_providers: ["bedrock"] + search_tool_name: agentcore-search diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec8a9336ca7..088d2193055 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3489,6 +3489,7 @@ class SearchProviders(str, Enum): YOU_COM = "you_com" APISERPENT = "apiserpent" TINYFISH = "tinyfish" + AGENTCORE = "agentcore" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 174bed09396..80a1f2b991f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8860,6 +8860,7 @@ class ProviderConfigManager: from litellm.llms.apiserpent.search.transformation import ( APISerpentSearchConfig, ) + from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.duckduckgo.search.transformation import DuckDuckGoSearchConfig @@ -8897,6 +8898,7 @@ class ProviderConfigManager: SearchProviders.YOU_COM: YouComSearchConfig, SearchProviders.APISERPENT: APISerpentSearchConfig, SearchProviders.TINYFISH: TinyfishSearchConfig, + SearchProviders.AGENTCORE: AgentCoreSearchConfig, } config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py new file mode 100644 index 00000000000..5d2e4d1c3cd --- /dev/null +++ b/tests/search_tests/test_agentcore_search.py @@ -0,0 +1,231 @@ +""" +Tests for Amazon Bedrock AgentCore Web Search integration. +""" + +import json +import os +import sys +import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig + +GATEWAY_URL = "https://testgateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + +MCP_RESULTS = [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "text": "Snippet for result 1", + "publishedDate": "2026-06-16", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "text": "Snippet for result 2", + }, +] + + +def _mcp_response_body() -> dict: + return { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(MCP_RESULTS)}]}, + } + + +def _make_mock_response(json_body: dict = None, text: str = None) -> MagicMock: + mock_response = MagicMock() + mock_response.status_code = 200 + if text is not None: + mock_response.text = text + else: + mock_response.text = json.dumps(json_body) + mock_response.json.return_value = json_body + return mock_response + + +class TestAgentCoreSearch: + """ + Tests for AgentCore Web Search functionality with mocked network/signing. + """ + + @pytest.mark.asyncio + async def test_agentcore_search_request_payload(self): + """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + + mock_response = _make_mock_response(_mcp_response_body()) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch.object( + AgentCoreSearchConfig, + "_sign_request", + return_value=( + {"Authorization": "AWS4-HMAC-SHA256 test", "Content-Type": "application/json"}, + json.dumps({"signed": True}).encode(), + ), + ) as mock_sign, + ): + mock_post.return_value = mock_response + + response = await litellm.asearch( + query="latest developments in AI", + search_provider="agentcore", + max_results=5, + ) + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"] == GATEWAY_URL + # Signed body must be sent verbatim + assert call_kwargs["data"] == json.dumps({"signed": True}).encode() + assert "json" not in call_kwargs + + # Signing was invoked with the MCP request + mock_sign.assert_called_once() + sign_kwargs = mock_sign.call_args.kwargs + request_data = sign_kwargs["request_data"] + assert request_data["method"] == "tools/call" + assert request_data["params"]["name"] == "web-search-tool___WebSearch" + assert request_data["params"]["arguments"]["query"] == "latest developments in AI" + assert request_data["params"]["arguments"]["maxResults"] == 5 + assert sign_kwargs["service_name"] == "bedrock-agentcore" + + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + assert response.results[0].url == "https://example.com/1" + assert response.results[0].snippet == "Snippet for result 1" + assert response.results[0].date == "2026-06-16" + + def test_transform_search_request_query_truncation(self): + """AgentCore rejects queries > 200 chars; the request must truncate.""" + config = AgentCoreSearchConfig() + long_query = "a" * 300 + data = config.transform_search_request(query=long_query, optional_params={}) + assert len(data["params"]["arguments"]["query"]) == 200 + + def test_transform_search_request_joins_list_queries(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query=["foo", "bar"], optional_params={}) + assert data["params"]["arguments"]["query"] == "foo bar" + + def test_transform_search_request_custom_tool_name(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) + assert data["params"]["name"] == "my-target___WebSearch" + + def test_get_complete_url_requires_gateway_url(self): + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + with pytest.raises(ValueError, match="AGENTCORE_GATEWAY_URL"): + config.get_complete_url(api_base=None, optional_params={}) + + def test_get_complete_url_prefers_api_base(self): + config = AgentCoreSearchConfig() + assert config.get_complete_url(api_base=GATEWAY_URL, optional_params={}) == GATEWAY_URL + + def test_validate_environment_sets_mcp_headers(self): + """MCP Streamable HTTP requires accepting both JSON and SSE.""" + config = AgentCoreSearchConfig() + headers = config.validate_environment(headers={}) + assert headers["Accept"] == "application/json, text/event-stream" + assert headers["Content-Type"] == "application/json" + + def test_transform_search_response_parses_sse_frame(self): + """Gateway may answer with an SSE-framed JSON-RPC message.""" + config = AgentCoreSearchConfig() + body = _mcp_response_body() + sse_text = f"event: message\ndata: {json.dumps(body)}\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[1].url == "https://example.com/2" + + def test_transform_search_response_raises_on_mcp_error(self): + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + {"jsonrpc": "2.0", "id": 1, "error": {"code": -32601, "message": "tool not found"}} + ) + with pytest.raises(Exception, match="tool not found"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_sign_request_uses_bearer_token_when_api_key_set(self): + """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" + config = AgentCoreSearchConfig() + request_data = {"jsonrpc": "2.0", "id": 1} + + headers, signed_body = config.sign_request( + headers={"Content-Type": "application/json"}, + optional_params={}, + request_data=request_data, + api_base=GATEWAY_URL, + api_key="test-jwt-token", + ) + assert headers["Authorization"] == "Bearer test-jwt-token" + assert signed_body == json.dumps(request_data).encode() + + def test_sign_request_uses_bearer_token_from_env(self): + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + + def test_sign_request_passes_explicit_aws_credentials(self): + """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIATEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + }, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + passed = mock_base_sign.call_args.kwargs["optional_params"] + assert passed["aws_access_key_id"] == "AKIATEST" + assert passed["aws_secret_access_key"] == "secret" + assert passed["aws_session_token"] == "token" + + def test_sign_request_derives_region_from_gateway_url(self): + """Signing region must come from the gateway URL, not the caller's default region.""" + config = AgentCoreSearchConfig() + eu_url = "https://gw-x.gateway.bedrock-agentcore.eu-central-1.amazonaws.com/mcp" + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=eu_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" From ebdad6e3ef5fa35daee25cd8bad9ae5b6e54f353 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Wed, 22 Jul 2026 02:26:25 +0000 Subject: [PATCH 02/88] fix: address bot review findings (auth hardening, SSE parsing, defaults) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Refuse to send the server-managed AGENTCORE_GATEWAY_TOKEN to a caller-supplied api_base (reuses resolve_server_api_key's trusted-host guard) — closes the token-exfiltration path via /search_tools/test_connection - Disable BaseAWSLLM's AWS_BEARER_TOKEN_BEDROCK fallback when signing: that token is a Bedrock Runtime credential and must not reach an AgentCore gateway - Parse SSE responses per spec: join multi-line data fields, iterate events, and return the JSON-RPC response (result/error) instead of the first data line — progress notifications no longer shadow the result - Validate tool_name ends with ___WebSearch so a caller-supplied name cannot invoke unrelated tools on the same gateway with the proxy's credentials - Send the documented maxResults default (10) explicitly instead of leaving it to the gateway - Custom gateway hostnames: raise a clear error when no signing region can be derived and none is configured, instead of signing for a guessed region - 7 new unit tests covering each fix (20 total) --- litellm/llms/bedrock/search/transformation.py | 93 ++++++++++++++++--- tests/search_tests/test_agentcore_search.py | 92 ++++++++++++++++++ 2 files changed, 170 insertions(+), 15 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 16d671f26b0..3bae84f0026 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -51,11 +51,20 @@ from litellm.secret_managers.main import get_secret_str # AgentCore web-search rejects queries longer than 200 characters AGENTCORE_MAX_QUERY_LENGTH = 200 +# The provider contract documents a default of 10 results — send it explicitly +# so the gateway can't silently apply a different default. +AGENTCORE_DEFAULT_MAX_RESULTS = 10 + # Default MCP tool name for a gateway web-search connector target: # "___". Override with AGENTCORE_SEARCH_TOOL_NAME # or optional_params["tool_name"] when the target uses a custom name. AGENTCORE_DEFAULT_TOOL_NAME = "web-search-tool___WebSearch" +# All web-search connector tools share this suffix; rejecting other names keeps +# a caller-supplied tool_name from invoking unrelated tools on the same gateway +# with the proxy's credentials. +AGENTCORE_TOOL_NAME_SUFFIX = "___WebSearch" + class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): def __init__(self) -> None: @@ -128,10 +137,15 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): or get_secret_str("AGENTCORE_SEARCH_TOOL_NAME") or AGENTCORE_DEFAULT_TOOL_NAME ) + if not tool_name.endswith(AGENTCORE_TOOL_NAME_SUFFIX): + raise ValueError( + f"Invalid AgentCore search tool_name '{tool_name}': must end with " + f"'{AGENTCORE_TOOL_NAME_SUFFIX}' (a web-search connector tool). " + "Other gateway tools cannot be invoked through this provider." + ) arguments: dict[str, Union[str, int]] = {"query": query} - if "max_results" in optional_params: - arguments["maxResults"] = optional_params["max_results"] + arguments["maxResults"] = optional_params.get("max_results", AGENTCORE_DEFAULT_MAX_RESULTS) return { "jsonrpc": "2.0", @@ -159,14 +173,28 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): if not isinstance(request_data, dict): raise ValueError("AgentCore search expects a single dict request body") - bearer_token = api_key or get_secret_str("AGENTCORE_GATEWAY_TOKEN") + # Server-managed token fallback is gated on the request targeting the + # operator-configured gateway host — otherwise an authenticated caller + # could point api_base at their own server (e.g. via + # /search_tools/test_connection) and receive AGENTCORE_GATEWAY_TOKEN. + bearer_token = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("AGENTCORE_GATEWAY_TOKEN",), + base_env_var="AGENTCORE_GATEWAY_URL", + default_api_base=None, + ) if bearer_token: headers["Authorization"] = f"Bearer {bearer_token}" return headers, json.dumps(request_data).encode() # The signing region must match the gateway's region — derive it from - # the gateway URL so callers don't have to set aws_region_name to a - # region different from their default. + # standard gateway hostnames so callers don't have to set + # aws_region_name to a region different from their default. Custom or + # private hostnames can't be parsed: fall back to an explicitly + # configured region (param or AWS env vars), and error out rather than + # silently signing for a guessed region the gateway would reject with + # a confusing auth error. signing_params = dict(optional_params) if signing_params.get("aws_region_name") is None: match = re.search( @@ -175,13 +203,23 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): ) if match: signing_params["aws_region_name"] = match.group(1) + elif not any(get_secret_str(var) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")): + raise ValueError( + f"Cannot derive the SigV4 signing region from api_base '{api_base}'. " + "Set aws_region_name (or the AWS_REGION env var) to the gateway's " + "region when using a custom hostname." + ) + # api_key="" (not None, but falsy) disables BaseAWSLLM's fallback to the + # AWS_BEARER_TOKEN_BEDROCK env var: that token is a Bedrock Runtime + # credential and must not be sent to an AgentCore gateway. return self._sign_request( service_name="bedrock-agentcore", headers=headers, optional_params=signing_params, request_data=request_data, api_base=api_base, + api_key="", ) def transform_search_response( @@ -231,17 +269,42 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): @staticmethod def _parse_mcp_body(raw_response: httpx.Response) -> dict: - """Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response.""" + """ + Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response. + + Per the SSE spec, an event's data is the concatenation of all its + ``data:`` lines (joined with newlines), and a stream may carry several + events (e.g. progress notifications before the JSON-RPC response). + Return the event whose payload carries the ``id``-matched JSON-RPC + response — i.e. one containing ``result`` or ``error``. + """ text = raw_response.text - if text.lstrip().startswith(("event:", "data:")): - for line in text.splitlines(): - if line.startswith("data:"): - return json.loads(line[len("data:") :].strip()) - raise BedrockError( - status_code=502, - message=f"AgentCore gateway returned SSE without a data frame: {text[:200]}", - ) - return raw_response.json() + if not text.lstrip().startswith(("event:", "data:", ":", "id:", "retry:")): + return raw_response.json() + + last_parsed: dict | None = None + data_lines: list[str] = [] + # Trailing sentinel flushes the final event even without a blank line + for line in text.splitlines() + [""]: + if line.startswith("data:"): + data_lines.append(line[len("data:") :].lstrip()) + continue + if line == "" and data_lines: + try: + parsed = json.loads("\n".join(data_lines)) + except json.JSONDecodeError: + parsed = None + data_lines = [] + if isinstance(parsed, dict): + last_parsed = parsed + if "result" in parsed or "error" in parsed: + return parsed + if last_parsed is not None: + return last_parsed + raise BedrockError( + status_code=502, + message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}", + ) def get_error_class( self, diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index 5d2e4d1c3cd..d5f6ab3e82f 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -123,6 +123,18 @@ class TestAgentCoreSearch: data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) assert data["params"]["name"] == "my-target___WebSearch" + def test_transform_search_request_rejects_non_websearch_tool_name(self): + """A caller-supplied tool_name must not reach other tools on the gateway.""" + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="must end with"): + config.transform_search_request(query="q", optional_params={"tool_name": "admin-target___DeleteUser"}) + + def test_transform_search_request_sends_documented_default_max_results(self): + """The documented default of 10 is sent explicitly, not left to the gateway.""" + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={}) + assert data["params"]["arguments"]["maxResults"] == 10 + def test_get_complete_url_requires_gateway_url(self): config = AgentCoreSearchConfig() os.environ.pop("AGENTCORE_GATEWAY_URL", None) @@ -151,6 +163,29 @@ class TestAgentCoreSearch: assert len(response.results) == 2 assert response.results[1].url == "https://example.com/2" + def test_transform_search_response_parses_multiline_sse_data(self): + """SSE data may be split across several data: lines (joined per spec).""" + config = AgentCoreSearchConfig() + pretty = json.dumps(_mcp_response_body(), indent=2) + sse_text = "event: message\n" + "\n".join(f"data: {line}" for line in pretty.splitlines()) + "\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_skips_progress_events(self): + """A progress notification before the JSON-RPC result must not shadow it.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\ndata: {json.dumps(progress)}\n\n" + f"event: message\ndata: {json.dumps(_mcp_response_body())}\n\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + def test_transform_search_response_raises_on_mcp_error(self): config = AgentCoreSearchConfig() mock_response = _make_mock_response( @@ -175,8 +210,10 @@ class TestAgentCoreSearch: assert signed_body == json.dumps(request_data).encode() def test_sign_request_uses_bearer_token_from_env(self): + """Server token is attached when the request targets the configured gateway host.""" config = AgentCoreSearchConfig() os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL try: headers, _ = config.sign_request( headers={}, @@ -187,6 +224,61 @@ class TestAgentCoreSearch: assert headers["Authorization"] == "Bearer env-jwt-token" finally: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_refuses_server_token_to_untrusted_host(self): + """Server-managed token must not be sent to a caller-chosen api_base.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://attacker.example.com/mcp", + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_does_not_leak_bedrock_bearer_token(self): + """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not + replace SigV4 on requests to an AgentCore gateway.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + # api_key="" (falsy, not None) disables the base class's + # AWS_BEARER_TOKEN_BEDROCK env fallback. + assert mock_base_sign.call_args.kwargs["api_key"] == "" + + def test_sign_request_custom_hostname_requires_region(self): + """Non-standard hostnames can't yield a signing region — require it explicitly.""" + config = AgentCoreSearchConfig() + saved = {var: os.environ.pop(var, None) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")} + try: + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + finally: + for var, val in saved.items(): + if val is not None: + os.environ[var] = val def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" From 2f342dc12d709666ca78f9d989313715d66ec216 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:27:10 +0000 Subject: [PATCH 03/88] fix: honor AWS shared-config region for custom gateway hostnames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous check only consulted AWS_REGION* env vars before rejecting custom hostnames, breaking deployments that configure their region via the AWS shared config (profile). Resolve through boto3's session (env vars + shared config) and only error when that chain yields nothing — never sign with a silently guessed region. --- litellm/llms/bedrock/search/transformation.py | 31 +++++++++++------ tests/search_tests/test_agentcore_search.py | 34 +++++++++++++++---- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 3bae84f0026..9ab651f952c 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -190,11 +190,11 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): # The signing region must match the gateway's region — derive it from # standard gateway hostnames so callers don't have to set - # aws_region_name to a region different from their default. Custom or - # private hostnames can't be parsed: fall back to an explicitly - # configured region (param or AWS env vars), and error out rather than - # silently signing for a guessed region the gateway would reject with - # a confusing auth error. + # aws_region_name to a region different from their default. For custom + # or private hostnames, defer to BaseAWSLLM's normal region resolution + # (params, env vars, AWS shared config / profile); only error out when + # that chain yields nothing, rather than silently signing for a guessed + # region the gateway would reject with a confusing auth error. signing_params = dict(optional_params) if signing_params.get("aws_region_name") is None: match = re.search( @@ -203,12 +203,21 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): ) if match: signing_params["aws_region_name"] = match.group(1) - elif not any(get_secret_str(var) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")): - raise ValueError( - f"Cannot derive the SigV4 signing region from api_base '{api_base}'. " - "Set aws_region_name (or the AWS_REGION env var) to the gateway's " - "region when using a custom hostname." - ) + else: + # boto3's session resolution covers env vars AND the AWS shared + # config (profile region) — unlike BaseAWSLLM's helper, which + # silently defaults to us-west-2 when nothing is configured. + import boto3 + + configured_region = boto3.Session().region_name + if configured_region: + signing_params["aws_region_name"] = configured_region + else: + raise ValueError( + f"Cannot derive the SigV4 signing region from api_base '{api_base}' " + "or the AWS configuration chain. Set aws_region_name (or AWS_REGION / " + "a profile region) to the gateway's region when using a custom hostname." + ) # api_key="" (not None, but falsy) disables BaseAWSLLM's fallback to the # AWS_BEARER_TOKEN_BEDROCK env var: that token is a Bedrock Runtime diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index d5f6ab3e82f..c041bfc8d50 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -264,10 +264,12 @@ class TestAgentCoreSearch: assert mock_base_sign.call_args.kwargs["api_key"] == "" def test_sign_request_custom_hostname_requires_region(self): - """Non-standard hostnames can't yield a signing region — require it explicitly.""" + """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() - saved = {var: os.environ.pop(var, None) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")} - try: + + mock_session = MagicMock() + mock_session.region_name = None # nothing configured anywhere + with patch("boto3.Session", return_value=mock_session): with pytest.raises(ValueError, match="signing region"): config.sign_request( headers={}, @@ -275,10 +277,28 @@ class TestAgentCoreSearch: request_data={"jsonrpc": "2.0"}, api_base="https://gateway.internal.example.com/mcp", ) - finally: - for var, val in saved.items(): - if val is not None: - os.environ[var] = val + + def test_sign_request_custom_hostname_uses_shared_config_region(self): + """Custom hostname + region from AWS shared config (profile) must be honored.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" From b43441814b8aaf51b3fab143424f6b89b80bf259 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Thu, 23 Jul 2026 15:31:10 +0000 Subject: [PATCH 04/88] test: mirror AgentCore search tests into tests/test_litellm for coverage Coverage collection runs against the sharded tests/test_litellm tree, so the provider tests living only in tests/search_tests were invisible to codecov (patch coverage reported ~31% despite the suite). Mirror them as tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py and add edge-case tests (malformed MCP content blocks, SSE without a JSON frame, notification-only streams, list request body, error-class mapping). transformation.py line coverage: 99% (26 tests x2 trees). --- tests/search_tests/test_agentcore_search.py | 55 +++ .../test_agentcore_search_transformation.py | 400 ++++++++++++++++++ 2 files changed, 455 insertions(+) create mode 100644 tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index c041bfc8d50..578d2b0f63e 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -341,3 +341,58 @@ class TestAgentCoreSearch: api_base=eu_url, ) assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" + + +class TestAgentCoreSearchEdgeCases: + """Branch coverage for response parsing and error mapping.""" + + def test_transform_search_response_skips_non_text_and_bad_json_blocks(self): + """Non-text blocks and unparseable text blocks are skipped, not fatal.""" + config = AgentCoreSearchConfig() + body = { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [ + {"type": "image", "data": "..."}, + {"type": "text", "text": "not-json"}, + {"type": "text", "text": json.dumps(["scalar", {"title": "T", "url": "u", "text": "s"}])}, + ] + }, + } + mock_response = _make_mock_response(body) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + # only the one dict item survives; non-dict list entries are skipped + assert len(response.results) == 1 + assert response.results[0].title == "T" + + def test_parse_mcp_body_sse_without_json_frame_raises(self): + """An SSE stream carrying no parseable JSON object is a 502.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response(text="event: ping\ndata: not-json\n\n") + with pytest.raises(Exception, match="SSE without a JSON data frame"): + config._parse_mcp_body(mock_response) + + def test_parse_mcp_body_returns_last_event_when_no_result_frame(self): + """A stream of only notifications returns the last parsed event.""" + config = AgentCoreSearchConfig() + note = {"jsonrpc": "2.0", "method": "notifications/progress"} + mock_response = _make_mock_response(text=f"data: {json.dumps(note)}\n\n") + assert config._parse_mcp_body(mock_response) == note + + def test_sign_request_rejects_list_request_body(self): + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="single dict"): + config.sign_request( + headers={}, + optional_params={}, + request_data=[{"jsonrpc": "2.0"}], + api_base=GATEWAY_URL, + ) + + def test_get_error_class_maps_status_and_message(self): + config = AgentCoreSearchConfig() + err = config.get_error_class(error_message="boom", status_code=503, headers={}) + assert getattr(err, "status_code", None) == 503 + assert "boom" in str(err) diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py new file mode 100644 index 00000000000..6bbaf66d3b3 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -0,0 +1,400 @@ +""" +Tests for Amazon Bedrock AgentCore Web Search integration. + +Mirror of tests/search_tests/test_agentcore_search.py placed in the +test_litellm tree so the AgentCoreSearchConfig transformation is exercised by +the sharded CI (coverage collection runs against this tree). +""" + +import json +import os + +import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +import litellm +from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig + +GATEWAY_URL = "https://testgateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + +MCP_RESULTS = [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "text": "Snippet for result 1", + "publishedDate": "2026-06-16", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "text": "Snippet for result 2", + }, +] + + +def _mcp_response_body() -> dict: + return { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(MCP_RESULTS)}]}, + } + + +def _make_mock_response(json_body: dict = None, text: str = None) -> MagicMock: + mock_response = MagicMock() + mock_response.status_code = 200 + if text is not None: + mock_response.text = text + else: + mock_response.text = json.dumps(json_body) + mock_response.json.return_value = json_body + return mock_response + + +class TestAgentCoreSearch: + """ + Tests for AgentCore Web Search functionality with mocked network/signing. + """ + + @pytest.mark.asyncio + async def test_agentcore_search_request_payload(self): + """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + + mock_response = _make_mock_response(_mcp_response_body()) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch.object( + AgentCoreSearchConfig, + "_sign_request", + return_value=( + {"Authorization": "AWS4-HMAC-SHA256 test", "Content-Type": "application/json"}, + json.dumps({"signed": True}).encode(), + ), + ) as mock_sign, + ): + mock_post.return_value = mock_response + + response = await litellm.asearch( + query="latest developments in AI", + search_provider="agentcore", + max_results=5, + ) + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"] == GATEWAY_URL + # Signed body must be sent verbatim + assert call_kwargs["data"] == json.dumps({"signed": True}).encode() + assert "json" not in call_kwargs + + # Signing was invoked with the MCP request + mock_sign.assert_called_once() + sign_kwargs = mock_sign.call_args.kwargs + request_data = sign_kwargs["request_data"] + assert request_data["method"] == "tools/call" + assert request_data["params"]["name"] == "web-search-tool___WebSearch" + assert request_data["params"]["arguments"]["query"] == "latest developments in AI" + assert request_data["params"]["arguments"]["maxResults"] == 5 + assert sign_kwargs["service_name"] == "bedrock-agentcore" + + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + assert response.results[0].url == "https://example.com/1" + assert response.results[0].snippet == "Snippet for result 1" + assert response.results[0].date == "2026-06-16" + + def test_transform_search_request_query_truncation(self): + """AgentCore rejects queries > 200 chars; the request must truncate.""" + config = AgentCoreSearchConfig() + long_query = "a" * 300 + data = config.transform_search_request(query=long_query, optional_params={}) + assert len(data["params"]["arguments"]["query"]) == 200 + + def test_transform_search_request_joins_list_queries(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query=["foo", "bar"], optional_params={}) + assert data["params"]["arguments"]["query"] == "foo bar" + + def test_transform_search_request_custom_tool_name(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) + assert data["params"]["name"] == "my-target___WebSearch" + + def test_transform_search_request_rejects_non_websearch_tool_name(self): + """A caller-supplied tool_name must not reach other tools on the gateway.""" + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="must end with"): + config.transform_search_request(query="q", optional_params={"tool_name": "admin-target___DeleteUser"}) + + def test_transform_search_request_sends_documented_default_max_results(self): + """The documented default of 10 is sent explicitly, not left to the gateway.""" + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={}) + assert data["params"]["arguments"]["maxResults"] == 10 + + def test_get_complete_url_requires_gateway_url(self): + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + with pytest.raises(ValueError, match="AGENTCORE_GATEWAY_URL"): + config.get_complete_url(api_base=None, optional_params={}) + + def test_get_complete_url_prefers_api_base(self): + config = AgentCoreSearchConfig() + assert config.get_complete_url(api_base=GATEWAY_URL, optional_params={}) == GATEWAY_URL + + def test_validate_environment_sets_mcp_headers(self): + """MCP Streamable HTTP requires accepting both JSON and SSE.""" + config = AgentCoreSearchConfig() + headers = config.validate_environment(headers={}) + assert headers["Accept"] == "application/json, text/event-stream" + assert headers["Content-Type"] == "application/json" + + def test_transform_search_response_parses_sse_frame(self): + """Gateway may answer with an SSE-framed JSON-RPC message.""" + config = AgentCoreSearchConfig() + body = _mcp_response_body() + sse_text = f"event: message\ndata: {json.dumps(body)}\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[1].url == "https://example.com/2" + + def test_transform_search_response_parses_multiline_sse_data(self): + """SSE data may be split across several data: lines (joined per spec).""" + config = AgentCoreSearchConfig() + pretty = json.dumps(_mcp_response_body(), indent=2) + sse_text = "event: message\n" + "\n".join(f"data: {line}" for line in pretty.splitlines()) + "\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_skips_progress_events(self): + """A progress notification before the JSON-RPC result must not shadow it.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\ndata: {json.dumps(progress)}\n\n" + f"event: message\ndata: {json.dumps(_mcp_response_body())}\n\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_raises_on_mcp_error(self): + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + {"jsonrpc": "2.0", "id": 1, "error": {"code": -32601, "message": "tool not found"}} + ) + with pytest.raises(Exception, match="tool not found"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_sign_request_uses_bearer_token_when_api_key_set(self): + """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" + config = AgentCoreSearchConfig() + request_data = {"jsonrpc": "2.0", "id": 1} + + headers, signed_body = config.sign_request( + headers={"Content-Type": "application/json"}, + optional_params={}, + request_data=request_data, + api_base=GATEWAY_URL, + api_key="test-jwt-token", + ) + assert headers["Authorization"] == "Bearer test-jwt-token" + assert signed_body == json.dumps(request_data).encode() + + def test_sign_request_uses_bearer_token_from_env(self): + """Server token is attached when the request targets the configured gateway host.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_refuses_server_token_to_untrusted_host(self): + """Server-managed token must not be sent to a caller-chosen api_base.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://attacker.example.com/mcp", + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_does_not_leak_bedrock_bearer_token(self): + """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not + replace SigV4 on requests to an AgentCore gateway.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + # api_key="" (falsy, not None) disables the base class's + # AWS_BEARER_TOKEN_BEDROCK env fallback. + assert mock_base_sign.call_args.kwargs["api_key"] == "" + + def test_sign_request_custom_hostname_requires_region(self): + """Custom hostname + empty AWS config chain → clear error, no guessed region.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = None # nothing configured anywhere + with patch("boto3.Session", return_value=mock_session): + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + + def test_sign_request_custom_hostname_uses_shared_config_region(self): + """Custom hostname + region from AWS shared config (profile) must be honored.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + + def test_sign_request_passes_explicit_aws_credentials(self): + """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIATEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + }, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + passed = mock_base_sign.call_args.kwargs["optional_params"] + assert passed["aws_access_key_id"] == "AKIATEST" + assert passed["aws_secret_access_key"] == "secret" + assert passed["aws_session_token"] == "token" + + def test_sign_request_derives_region_from_gateway_url(self): + """Signing region must come from the gateway URL, not the caller's default region.""" + config = AgentCoreSearchConfig() + eu_url = "https://gw-x.gateway.bedrock-agentcore.eu-central-1.amazonaws.com/mcp" + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=eu_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" + + +class TestAgentCoreSearchEdgeCases: + """Branch coverage for response parsing and error mapping.""" + + def test_transform_search_response_skips_non_text_and_bad_json_blocks(self): + """Non-text blocks and unparseable text blocks are skipped, not fatal.""" + config = AgentCoreSearchConfig() + body = { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [ + {"type": "image", "data": "..."}, + {"type": "text", "text": "not-json"}, + {"type": "text", "text": json.dumps(["scalar", {"title": "T", "url": "u", "text": "s"}])}, + ] + }, + } + mock_response = _make_mock_response(body) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + # only the one dict item survives; non-dict list entries are skipped + assert len(response.results) == 1 + assert response.results[0].title == "T" + + def test_parse_mcp_body_sse_without_json_frame_raises(self): + """An SSE stream carrying no parseable JSON object is a 502.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response(text="event: ping\ndata: not-json\n\n") + with pytest.raises(Exception, match="SSE without a JSON data frame"): + config._parse_mcp_body(mock_response) + + def test_parse_mcp_body_returns_last_event_when_no_result_frame(self): + """A stream of only notifications returns the last parsed event.""" + config = AgentCoreSearchConfig() + note = {"jsonrpc": "2.0", "method": "notifications/progress"} + mock_response = _make_mock_response(text=f"data: {json.dumps(note)}\n\n") + assert config._parse_mcp_body(mock_response) == note + + def test_sign_request_rejects_list_request_body(self): + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="single dict"): + config.sign_request( + headers={}, + optional_params={}, + request_data=[{"jsonrpc": "2.0"}], + api_base=GATEWAY_URL, + ) + + def test_get_error_class_maps_status_and_message(self): + config = AgentCoreSearchConfig() + err = config.get_error_class(error_message="boom", status_code=503, headers={}) + assert getattr(err, "status_code", None) == 503 + assert "boom" in str(err) From e1629b77dbce0e9eaddbdb726e13ed7edd614ba9 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Sun, 26 Jul 2026 07:26:46 +0000 Subject: [PATCH 05/88] fix(interactions): sync queued status enum from #34318 to unblock CI on stale daily branch --- litellm/types/interactions/generated.py | 2 ++ tests/test_litellm/interactions/test_openapi_compliance.py | 1 + 2 files changed, 3 insertions(+) diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 793cc02ff17..4a1ef5ed696 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -173,6 +173,7 @@ class Status1(Enum): cancelled = "cancelled" incomplete = "incomplete" budget_exceeded = "budget_exceeded" + queued = "queued" class InteractionStatusUpdate(BaseModel): @@ -341,6 +342,7 @@ class Status3(Enum): CANCELLED = "cancelled" INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" + QUEUED = "queued" class ModelOption(RootModel[str]): diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 209e99895db..11b08fa45a8 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -194,6 +194,7 @@ class TestResponseCompliance: "cancelled", "incomplete", "budget_exceeded", + "queued", ] assert status_prop["enum"] == expected_statuses print(f"✓ Status enum values: {expected_statuses}") From af2246c5b8d75bcaefb67d5615063183bf5e7502 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:47:05 +0000 Subject: [PATCH 06/88] fix(anthropic,bedrock): report provider thinking tokens instead of classifying them as text Resolves LIT-5244 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 2 +- litellm/llms/anthropic/chat/transformation.py | 77 ++++++++++-- .../bedrock/chat/converse_transformation.py | 19 ++- litellm/llms/bedrock/chat/invoke_handler.py | 8 +- .../transformation.py | 34 ++++-- litellm/types/llms/anthropic.py | 6 + .../test_streaming_chunk_builder_utils.py | 46 +++++++ .../test_anthropic_chat_transformation.py | 112 ++++++++++++++++++ .../chat/test_converse_transformation.py | 81 +++++++++++++ .../test_reasoning_content_transformation.py | 101 ++++++++++++++++ 10 files changed, 462 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index a2f9c80f577..fe51b5cc822 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -583,7 +583,7 @@ class ChunkProcessor: for choice in response.choices: if ( hasattr(cast(Choices, choice).message, "reasoning_content") - and cast(Choices, choice).message.reasoning_content is not None + and cast(Choices, choice).message.reasoning_content ): if reasoning_tokens is None: reasoning_tokens = 0 diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1f9022bf28f..5c27535014a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1,9 +1,11 @@ import json import re import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, NoReturn, cast import httpx +from pydantic import ValidationError import litellm from litellm.constants import ( @@ -38,6 +40,7 @@ from litellm.types.llms.anthropic import ( AnthropicMessagesTool, AnthropicMessagesToolChoice, AnthropicOutputSchema, + AnthropicOutputTokensDetails, AnthropicSystemMessageContent, AnthropicThinkingParam, AnthropicWebSearchTool, @@ -2104,6 +2107,66 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): compaction_blocks, ) + @staticmethod + def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + details: Final = usage_object.get("output_tokens_details") + if not isinstance(details, Mapping): + return None + try: + return AnthropicOutputTokensDetails.model_validate(details).thinking_tokens + except ValidationError: + return None + + @staticmethod + def _response_has_thinking_block(completion_response: Mapping[str, object] | None) -> bool: + if completion_response is None: + return False + content: Final = completion_response.get("content") + if not isinstance(content, list): + return False + return any( + isinstance(block, Mapping) and block.get("type") in ("thinking", "redacted_thinking") for block in content + ) + + def _build_completion_token_details( + self, + usage_object: Mapping[str, object], + iterations: Sequence[object] | None, + completion_tokens: int, + reasoning_content: str | None, + completion_response: Mapping[str, object] | None, + ) -> CompletionTokensDetailsWrapper: + reported_thinking_tokens: Final = ( + self._sum_iteration_thinking_tokens(iterations) + if iterations + else self._thinking_tokens_from_usage(usage_object) + ) + if reported_thinking_tokens is not None: + capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) + return CompletionTokensDetailsWrapper( + reasoning_tokens=capped_reported, + text_tokens=completion_tokens - capped_reported, + ) + if reasoning_content: + estimated: Final = min( + token_counter(text=reasoning_content, count_response_tokens=True), + completion_tokens, + ) + return CompletionTokensDetailsWrapper( + reasoning_tokens=max(0, estimated), + text_tokens=completion_tokens - max(0, estimated), + ) + if self._response_has_thinking_block(completion_response): + return CompletionTokensDetailsWrapper(reasoning_tokens=None, text_tokens=None) + return CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=completion_tokens) + + def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: + per_iteration: Final = tuple( + self._thinking_tokens_from_usage(iteration) for iteration in iterations if isinstance(iteration, Mapping) + ) + reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) + return sum(reported) if reported else None + def calculate_usage( self, usage_object: dict, @@ -2182,14 +2245,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): cache_creation_token_details=cache_creation_token_details, text_tokens=raw_input_tokens, ) - # Always populate completion_token_details, not just when there's reasoning_content - estimated_reasoning_tokens: Final = ( - token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 - ) - reasoning_tokens: Final = min(estimated_reasoning_tokens, completion_tokens) - completion_token_details: Final = CompletionTokensDetailsWrapper( - reasoning_tokens=max(0, reasoning_tokens), - text_tokens=(completion_tokens - reasoning_tokens if reasoning_tokens > 0 else completion_tokens), + completion_token_details: Final = self._build_completion_token_details( + usage_object=_usage, + iterations=iterations, + completion_tokens=completion_tokens, + reasoning_content=reasoning_content, + completion_response=completion_response, ) total_tokens: Final = prompt_tokens + completion_tokens diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 91adff50a17..93feabe7222 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1764,6 +1764,7 @@ class AmazonConverseConfig(BaseConfig): self, usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, + thinking_ran: bool = False, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1784,10 +1785,19 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 - completion_tokens_details: Final = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens, - text_tokens=(output_tokens - reasoning_tokens if reasoning_tokens > 0 else output_tokens), + reasoning_tokens: Final = ( + token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 + ) + completion_tokens_details: Final = ( + CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens, + text_tokens=output_tokens - reasoning_tokens, + ) + if reasoning_tokens > 0 + else CompletionTokensDetailsWrapper( + reasoning_tokens=None if thinking_ran else 0, + text_tokens=None if thinking_ran else output_tokens, + ) ) openai_usage: Final = Usage( prompt_tokens=input_tokens, @@ -2184,6 +2194,7 @@ class AmazonConverseConfig(BaseConfig): usage: Final = self._transform_usage( completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), + thinking_ran=reasoningContentBlocks is not None, ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index a2bb179f72f..57510ff334d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -330,6 +330,7 @@ class AWSEventStreamDecoder: self.response_id: str | None = None self.json_mode = json_mode self._current_tool_name: str | None = None + self._thinking_ran = False def check_empty_tool_call_args(self) -> bool: """ @@ -559,7 +560,12 @@ class AWSEventStreamDecoder: elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) elif "usage" in chunk_data: - usage = converse_config._transform_usage(chunk_data.get("usage", {})) + usage = converse_config._transform_usage( + chunk_data.get("usage", {}), + thinking_ran=self._thinking_ran, + ) + if thinking_blocks: + self._thinking_ran = True model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 79e05545358..174a55aac85 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,7 +4,7 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast from openai.types.chat.chat_completion_named_tool_choice_param import ( @@ -1745,6 +1745,12 @@ class LiteLLMCompletionResponsesConfig: output_items.append(item) return output_items + @staticmethod + def _encode_thinking_blocks(message: Message) -> str | None: + thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () + preserved: Final = tuple(block for block in thinking_blocks if block.get("signature") or block.get("data")) + return json.dumps(preserved, separators=(",", ":")) if preserved else None + @staticmethod def _extract_reasoning_output_items( chat_completion_response: ModelResponse, @@ -1753,23 +1759,31 @@ class LiteLLMCompletionResponsesConfig: for choice in choices: if hasattr(choice, "message") and choice.message: message = choice.message - if hasattr(message, "reasoning_content") and message.reasoning_content: + reasoning_content = getattr(message, "reasoning_content", None) or "" + encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message) + if reasoning_content or encrypted_content: # Only check the first choice for reasoning content return [ GenericResponseOutputItem( type="reasoning", - id=f"rs_{hash(str(message.reasoning_content))}", + id=f"rs_{hash(reasoning_content or encrypted_content)}", status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( choice.finish_reason ), role="assistant", - content=[ - OutputText( - type="output_text", - text=message.reasoning_content, - annotations=[], - ) - ], + content=( + [ + OutputText( + type="output_text", + text=reasoning_content, + annotations=[], + ) + ] + if reasoning_content + # mutable-ok: GenericResponseOutputItem.content is typed as a list + else [] + ), + encrypted_content=encrypted_content, ) ] return [] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 95f8db66eda..f111d3c6e56 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -626,6 +626,12 @@ class AnthropicResponseUsageBlock(BaseModel): output_tokens: int +class AnthropicOutputTokensDetails(BaseModel): + model_config = ConfigDict(extra="allow") + + thinking_tokens: Optional[int] = None + + AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 0114db381cf..15bbe476a06 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1180,3 +1180,49 @@ def test_get_combined_tool_content_joins_many_custom_tool_input_fragments_in_ord assert isinstance(combined[1], ChatCompletionMessageCustomToolCall) assert combined[1].custom.name == "run_script" assert combined[1].custom.input == "".join(object_fragments) + + +def _reasoning_stream_chunk() -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-reasoning", + model="claude-opus-4-8", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="10", role="assistant"))], + ) + + +def test_count_reasoning_tokens_returns_none_for_signature_only_thinking(): + from litellm.types.utils import Choices, Message, ModelResponse + + processor = ChunkProcessor(chunks=[_reasoning_stream_chunk()]) + response = ModelResponse( + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="10", role="assistant", reasoning_content=""), + ) + ] + ) + + assert processor.count_reasoning_tokens(response) is None + + +def test_count_reasoning_tokens_counts_visible_reasoning(): + from litellm.types.utils import Choices, Message, ModelResponse + + processor = ChunkProcessor(chunks=[_reasoning_stream_chunk()]) + response = ModelResponse( + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="let me count the primes under thirty", + ), + ) + ] + ) + + assert processor.count_reasoning_tokens(response) > 0 diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 94a4a3fc945..063b965dd47 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -119,6 +119,118 @@ def test_calculate_usage_clamps_text_tokens_when_reasoning_estimate_exceeds_outp assert usage.completion_tokens_details.text_tokens == 0 +def test_calculate_usage_prefers_provider_reported_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 32, + "output_tokens": 421, + "output_tokens_details": {"thinking_tokens": 372}, + }, + reasoning_content="", + completion_response={ + "content": [ + {"type": "thinking", "thinking": "", "signature": "sig"}, + {"type": "text", "text": "10"}, + ] + }, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 372 + assert usage.completion_tokens_details.text_tokens == 49 + + +def test_calculate_usage_provider_thinking_tokens_win_over_visible_reasoning_estimate(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 50, + "output_tokens": 811, + "output_tokens_details": {"thinking_tokens": 747}, + }, + reasoning_content="short visible reasoning that tokenizes to far fewer than 747 tokens", + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 747 + assert usage.completion_tokens_details.text_tokens == 64 + + +def test_calculate_usage_sums_provider_thinking_tokens_across_iterations(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200, "output_tokens_details": {"thinking_tokens": 90}}, + ], + }, + reasoning_content=None, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 150 + assert usage.completion_tokens_details.text_tokens == 150 + + +def test_calculate_usage_reports_unknown_split_when_thinking_ran_without_a_count(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 32, "output_tokens": 580}, + reasoning_content="", + completion_response={ + "content": [ + {"type": "redacted_thinking", "data": "encrypted"}, + {"type": "text", "text": "10"}, + ] + }, + ) + + assert usage.completion_tokens == 580 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + +def test_calculate_usage_without_thinking_reports_all_output_as_text(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 32, "output_tokens": 171}, + reasoning_content=None, + completion_response={"content": [{"type": "text", "text": "10"}]}, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 171 + + +def test_calculate_usage_ignores_malformed_provider_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 32, + "output_tokens": 100, + "output_tokens_details": {"thinking_tokens": "not-a-number"}, + }, + reasoning_content=None, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 100 + + def test_calculate_usage_handles_mocked_output_tokens_with_reasoning_content(): config = AnthropicConfig() diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 6d318bb8729..1f759b58cf7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5934,3 +5934,84 @@ def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): ) assert "thinking" not in optional_params + + +def test_converse_usage_reports_unknown_split_for_signature_only_thinking(): + config = AmazonConverseConfig() + + usage = config._transform_usage( + ConverseTokenUsageBlock(inputTokens=32, outputTokens=581, totalTokens=613), + reasoning_content="", + thinking_ran=True, + ) + + assert usage.completion_tokens == 581 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + +def test_converse_usage_estimates_split_for_visible_thinking(): + config = AmazonConverseConfig() + + usage = config._transform_usage( + ConverseTokenUsageBlock(inputTokens=32, outputTokens=581, totalTokens=613), + reasoning_content="Let me think about how many primes there are under thirty.", + thinking_ran=True, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert ( + usage.completion_tokens_details.reasoning_tokens + usage.completion_tokens_details.text_tokens + == usage.completion_tokens + ) + + +def test_converse_usage_without_thinking_reports_all_output_as_text(): + config = AmazonConverseConfig() + + usage = config._transform_usage(ConverseTokenUsageBlock(inputTokens=32, outputTokens=171, totalTokens=203)) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 171 + + +def test_converse_transform_response_signature_only_thinking_reports_unknown_split(): + config = AmazonConverseConfig() + raw_response = MagicMock(status_code=200) + raw_response.text = json.dumps( + { + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "", "signature": "sig"}}}, + {"text": "10"}, + ], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 32, "outputTokens": 581, "totalTokens": 613}, + } + ) + raw_response.json.return_value = json.loads(raw_response.text) + + response = config._transform_response( + model="bedrock/global.anthropic.claude-opus-4-8", + response=raw_response, + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data={}, + messages=[], + encoding=None, + ) + + assert response.choices[0].message.reasoning_content == "" + + assert response.usage.completion_tokens_details.reasoning_tokens is None + assert response.usage.completion_tokens_details.text_tokens is None diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py index 020b5de0a2a..3c1980152a7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py @@ -263,6 +263,107 @@ class TestReasoningContentFinalResponse: assert len(reasoning_items) == 1, "Should have exactly one reasoning item" assert reasoning_items[0].content[0].text == "Reasoning for first answer" + def test_signature_only_thinking_block_still_emits_reasoning_item(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="", + thinking_blocks=[ + {"type": "thinking", "thinking": "", "signature": "signature-payload"} + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1, "Signature-only thinking should still surface a reasoning item" + assert reasoning_items[0].content == [] + assert "signature-payload" in reasoning_items[0].encrypted_content + + def test_redacted_thinking_block_preserved_as_encrypted_content(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + thinking_blocks=[{"type": "redacted_thinking", "data": "redacted-payload"}], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1 + assert "redacted-payload" in reasoning_items[0].encrypted_content + + def test_visible_thinking_keeps_text_and_signature(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="counting the primes", + thinking_blocks=[ + {"type": "thinking", "thinking": "counting the primes", "signature": "sig"} + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1 + assert reasoning_items[0].content[0].text == "counting the primes" + assert "sig" in reasoning_items[0].encrypted_content + def test_streaming_chunk_id_raw(): """Test that streaming chunk IDs are raw (not encoded) to match OpenAI format""" From 53ee9c8293d6d1aeb038eb1a674e5d8ad090dbef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:34:36 +0000 Subject: [PATCH 07/88] fix(anthropic): fall back when only some compaction iterations report thinking tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 10 +++-- .../transformation.py | 21 ++++----- litellm/types/llms/anthropic.py | 2 +- .../test_anthropic_chat_transformation.py | 44 +++++++++++++++++++ 4 files changed, 60 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e0e11be356c..feb26b19981 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2134,9 +2134,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reasoning_content: str | None, completion_response: Mapping[str, object] | None, ) -> CompletionTokensDetailsWrapper: + iteration_thinking_tokens: Final = self._sum_iteration_thinking_tokens(iterations) if iterations else None reported_thinking_tokens: Final = ( - self._sum_iteration_thinking_tokens(iterations) - if iterations + iteration_thinking_tokens + if iteration_thinking_tokens is not None else self._thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: @@ -2160,10 +2161,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) for iteration in iterations if isinstance(iteration, Mapping) + self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) - return sum(reported) if reported else None + return sum(reported) if len(reported) == len(per_iteration) else None @staticmethod def is_anthropic_usage_object(usage_object: dict) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fa2ce0d1505..f0614f1cacf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1763,18 +1763,15 @@ class LiteLLMCompletionResponsesConfig: choice.finish_reason ), role="assistant", - content=( - [ - OutputText( - type="output_text", - text=reasoning_content, - annotations=[], - ) - ] - if reasoning_content - # mutable-ok: GenericResponseOutputItem.content is typed as a list - else [] - ), + content=[ + OutputText( + type="output_text", + text=text, + annotations=[], + ) + for text in (reasoning_content,) + if text + ], encrypted_content=encrypted_content, ) ] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 7de383f6f13..f6b256ad5df 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -612,7 +612,7 @@ class AnthropicResponseUsageBlock(BaseModel): class AnthropicOutputTokensDetails(BaseModel): model_config = ConfigDict(extra="allow") - thinking_tokens: Optional[int] = None + thinking_tokens: int | None = None AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 828a9c30fb9..de62c990f11 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -180,6 +180,50 @@ def test_calculate_usage_sums_provider_thinking_tokens_across_iterations(): assert usage.completion_tokens_details.text_tokens == 150 +def test_calculate_usage_falls_back_when_only_some_iterations_report_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "output_tokens_details": {"thinking_tokens": 240}, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200}, + ], + }, + reasoning_content=None, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 240 + assert usage.completion_tokens_details.text_tokens == 60 + + +def test_calculate_usage_reports_unknown_split_when_only_some_iterations_report_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200}, + ], + }, + reasoning_content="", + completion_response={"content": [{"type": "thinking", "thinking": "", "signature": "sig"}]}, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + def test_calculate_usage_reports_unknown_split_when_thinking_ran_without_a_count(): config = AnthropicConfig() From b1d77bb5dbc4e2697783258385eb9ac8029fe5b4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 01:36:33 +0000 Subject: [PATCH 08/88] style: ruff format agentcore search transformation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 4faaad6cc00..f6d852ffb0a 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -118,9 +118,7 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: progress notifications before the JSON-RPC response. """ for chunk in _SSE_EVENT_SEPARATOR.split(text): - payload = "\n".join( - line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:") - ) + payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:")) if not payload: continue try: From 15a6664171f2ba2b559db1ea333db44996a9a7df Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:17:44 +0000 Subject: [PATCH 09/88] fix(search): keep signed auth headers out of logging callbacks Log the pre-signing headers in the search pre_call hook so SigV4 and bearer Authorization values are never handed to user-configured logger callbacks, and tighten sign_request's annotations. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/base_llm/search/transformation.py | 8 ++++---- litellm/llms/bedrock/search/transformation.py | 8 ++++---- litellm/llms/custom_httpx/llm_http_handler.py | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index fad4be538c7..59039d68ede 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -180,12 +180,12 @@ class BaseSearchConfig: def sign_request( self, - headers: dict, # mutable-ok: matches the request header dict every other hook on this base takes - optional_params: dict, # mutable-ok: matches the optional params dict every other hook on this base takes - request_data: dict | list[dict], # mutable-ok: matches transform_search_request's JSON body return type + headers: dict[str, str], # mutable-ok: matches the request header dict every other hook on this base takes + optional_params: dict[str, object], # mutable-ok: matches every other hook on this base + request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: transform_search_request's body api_base: str, api_key: str | None = None, - ) -> tuple[dict, bytes | None]: # mutable-ok: the handler passes these headers straight to httpx + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the handler passes these headers straight to httpx """ OPTIONAL diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index f6d852ffb0a..ca9759ed151 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -221,12 +221,12 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): def sign_request( self, - headers: dict, # mutable-ok: BaseSearchConfig hands providers the mutable request header dict - optional_params: dict, # mutable-ok: BaseSearchConfig passes optional params as a dict - request_data: dict | list[dict], # mutable-ok: BaseSearchConfig request bodies are JSON dicts + headers: dict[str, str], # mutable-ok: BaseSearchConfig hands providers the mutable request header dict + optional_params: dict[str, object], # mutable-ok: BaseSearchConfig passes optional params as a dict + request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: request bodies are JSON dicts api_base: str, api_key: str | None = None, - ) -> tuple[dict, bytes | None]: # mutable-ok: BaseSearchConfig.sign_request returns httpx headers + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseSearchConfig.sign_request returns httpx headers """ Authenticate the MCP request. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3ee79241d30..b2f8c21d83c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1755,7 +1755,7 @@ class BaseLLMHTTPHandler: additional_args={ "complete_input_dict": data, "api_base": complete_url, - "headers": signed_headers, + "headers": headers, }, ) @@ -1848,7 +1848,7 @@ class BaseLLMHTTPHandler: additional_args={ "complete_input_dict": data, "api_base": complete_url, - "headers": signed_headers, + "headers": headers, }, ) From 25144fc03cebf686483acdae02bf3bd1ce72e3d3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:11:12 +0000 Subject: [PATCH 10/88] chore: retrigger ci after docs merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 90cd378a5943bfa739b5eb2d1f4c4234a76920ac Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 01:04:52 +0000 Subject: [PATCH 11/88] fix(streaming): accept provider cost objects when propagating usage cost Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 19 ++++- .../test_streaming_handler.py | 74 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..48e58c578e6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1789,6 +1789,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = "tool_calls" return model_response + @staticmethod + def _resolve_provider_reported_cost(usage_cost: object) -> float | None: + """ + Providers report usage.cost either as a number or, for Perplexity, as a + breakdown object whose total lives under ``total_cost``. + """ + if isinstance(usage_cost, bool): + return None + if isinstance(usage_cost, (int, float)): + return float(usage_cost) + if isinstance(usage_cost, dict): + return CustomStreamWrapper._resolve_provider_reported_cost(usage_cost.get("total_cost")) + return None + @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", @@ -1799,10 +1813,11 @@ class CustomStreamWrapper: calculator uses it instead of a token-based estimate. """ _usage: Final[Usage | None] = getattr(response, "usage", None) - if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: + _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) + if _cost is not None: if "additional_headers" not in response._hidden_params: response._hidden_params["additional_headers"] = {} - response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost) + response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost def __next__(self) -> "ModelResponseStream": cache_hit = False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 101935cac0a..dad4faa98b4 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1676,6 +1676,80 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert provider_cost == 0.00025 +def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): + """ + Regression: Perplexity reports usage.cost as a breakdown object, which used to + blow up the end of the stream with + `float() argument must be a string or a real number, not 'dict'`. + """ + import litellm + from litellm.cost_calculator import get_response_cost_from_hidden_params + + chunks = [ + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056047, + model="perplexity/sonar", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hi", role="assistant"), + ) + ], + usage=None, + ), + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056048, + model="perplexity/sonar", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=None, + ), + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056049, + model="perplexity/sonar", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) + ], + usage=Usage( + completion_tokens=18, + prompt_tokens=12, + total_tokens=30, + cost={ + "input_tokens_cost": 0.000012, + "output_tokens_cost": 0.000018, + "request_cost": 0.005, + "total_cost": 0.00503, + }, + ), + ), + ] + + complete_response = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "test"}] + ) + + assert complete_response is not None + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + + assert ( + get_response_cost_from_hidden_params(complete_response._hidden_params) + == 0.00503 + ) + + +def test_provider_reported_cost_ignores_unusable_shapes(): + assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None + assert CustomStreamWrapper._resolve_provider_reported_cost({}) is None + assert CustomStreamWrapper._resolve_provider_reported_cost({"total_cost": None}) is None + assert CustomStreamWrapper._resolve_provider_reported_cost(0.5) == 0.5 + + def test_handle_special_delta_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 79a6d2b8d264502040698ff1856671669d4b5795 Mon Sep 17 00:00:00 2001 From: RayJueWang <570828708@qq.com> Date: Tue, 28 Jul 2026 14:01:34 +0800 Subject: [PATCH 12/88] fix(proxy): retry spend updates on Postgres deadlock instead of dropping them Spend-update transactions increment non-idempotent counters (spend = spend + x) inside prisma interactive transactions. Every retry loop only caught DB_RETRY_SAFE_ERROR_TYPES (httpx.ConnectError); a Postgres deadlock (SQLSTATE 40P01, surfaced by prisma as transaction conflict code P2034) fell through to a bare except that re-raised immediately, so on multi-pod / high-concurrency deployments any pod that lost a deadlock silently dropped its increment. A deadlock is replay-safe even though the increment is non-idempotent: Postgres aborts and fully rolls back the victim transaction, so no partial spend is committed. Add PrismaDBExceptionHandler.is_deadlock_error and route every spend path (user, end-user/key, team, team_member, org, tag/agent via _update_entity_spend_in_db, and the daily-spend upsert) through a shared _handle_spend_update_failure that retries connection errors and deadlocks with randomized jitter backoff and re-raises everything else or on exhaustion. --- litellm/proxy/db/db_spend_update_writer.py | 143 ++++++-------- litellm/proxy/db/exception_handler.py | 12 ++ .../proxy/db/test_db_spend_update_writer.py | 184 ++++++++++++++++++ .../proxy/db/test_exception_handler.py | 33 ++++ 4 files changed, 293 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b2b72c1cac4..16389cda336 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1122,6 +1122,23 @@ class DBSpendUpdateWriter: except Exception as e: verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) + @staticmethod + async def _handle_spend_update_failure( + e: Exception, + attempt: int, + n_retry_times: int, + start_time: float, + proxy_logging_obj: ProxyLogging, + ) -> None: + """Retry a failed spend-update transaction on connection errors or deadlocks, else re-raise.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + from litellm.proxy.utils import _raise_failed_update_spend_exception + + is_retryable = isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) or PrismaDBExceptionHandler.is_deadlock_error(e) + if not is_retryable or attempt >= n_retry_times: + _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + await asyncio.sleep(random.uniform(2**attempt, 2 ** (attempt + 1))) + async def _commit_spend_updates_to_db( self, prisma_client: PrismaClient, @@ -1133,10 +1150,7 @@ class DBSpendUpdateWriter: Commits all the spend `UPDATE` transactions to the Database """ - from litellm.proxy.utils import ( - ProxyUpdateSpend, - _raise_failed_update_spend_exception, - ) + from litellm.proxy.utils import ProxyUpdateSpend ### UPDATE USER TABLE ### user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"] @@ -1156,18 +1170,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE END-USER TABLE ### @@ -1199,18 +1208,13 @@ class DBSpendUpdateWriter: }, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TEAM TABLE ### @@ -1232,18 +1236,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TEAM Membership TABLE with spend ### @@ -1279,18 +1278,13 @@ class DBSpendUpdateWriter: ) # Transaction succeeded, break out of retry loop break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) # Invalidate cache for updated team memberships @@ -1321,25 +1315,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep( - # Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are - # cancelled basically at the same time, so if they wait the same time they will also retry at the same time - # and thus they are more likely to deadlock again. - # Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of - # repeated deadlocks, and therefore of exceeding the retry limit. - random.uniform(2**i, 2 ** (i + 1)) - ) except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TAG TABLE ### @@ -1388,8 +1370,6 @@ class DBSpendUpdateWriter: prisma_client: Prisma client instance proxy_logging_obj: Proxy logging object """ - from litellm.proxy.utils import _raise_failed_update_spend_exception - verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions) if transactions is not None and len(transactions.keys()) > 0: for i in range(n_retry_times + 1): @@ -1411,17 +1391,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await DBSpendUpdateWriter._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) # fmt: off @@ -1590,7 +1566,16 @@ class DBSpendUpdateWriter: break - except DB_RETRY_SAFE_ERROR_TYPES as e: + except Exception as e: + from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + ) + + is_retryable = isinstance( + e, DB_RETRY_SAFE_ERROR_TYPES + ) or PrismaDBExceptionHandler.is_deadlock_error(e) + if not is_retryable: + raise if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index e0a21ceed26..91c7e576dff 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -166,6 +166,18 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_deadlock_error(e: Exception) -> bool: + """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" + import prisma + + if not isinstance(e, prisma.errors.PrismaError): + return False + if getattr(e, "code", None) == "P2034": + return True + error_message = str(e).lower() + return "deadlock detected" in error_message or "40p01" in error_message + @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: """True iff ``e`` is a non-``PrismaError`` exception raised from inside diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index ca7d5fcd273..49ef653d4e1 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2268,3 +2268,187 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count assert internal["autorouter_savings_spend"] == 0.0 assert user_sent["api_requests"] == 1 assert user_sent["successful_requests"] == 1 + + +def _deadlock_error(): + from prisma.errors import RawQueryError + + return RawQueryError( + data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "LiteLLM_VerificationToken"}}} + ) + + +def _empty_spend_transactions(**overrides): + base = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + return {**base, **overrides} + + +def _good_tx(mock_batcher): + tx = AsyncMock() + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=False) + tx.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + return tx + + +def _failing_tx(error): + tx = MagicMock() + tx.__aenter__ = AsyncMock(side_effect=error) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +@pytest.mark.asyncio +async def test_commit_spend_updates_retries_deadlock_then_commits(monkeypatch): + """Regression: a deadlock on the key-spend UPDATE is retried and commits the increment exactly once.""" + slept = [] + monkeypatch.setattr( + "litellm.proxy.db.db_spend_update_writer.asyncio.sleep", + AsyncMock(side_effect=lambda s: slept.append(s)), + ) + + mock_batcher = MagicMock() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + assert mock_prisma_client.db.tx.call_count == 2 + mock_batcher.litellm_verificationtoken.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + assert call_kwargs["where"] == {"token": "sk-abc"} + assert call_kwargs["data"]["spend"] == {"increment": 0.5} + assert len(slept) == 1 + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_commit_spend_updates_raises_after_exhausting_deadlock_retries(monkeypatch): + """A deadlock that never clears must surface after the retry budget is spent, not loop or swallow.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(_deadlock_error())) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + from prisma.errors import RawQueryError + + with pytest.raises(RawQueryError): + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=2, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + assert mock_prisma_client.db.tx.call_count == 3 + + +@pytest.mark.asyncio +async def test_commit_spend_updates_does_not_retry_non_deadlock_data_error(monkeypatch): + """A non-retryable data-layer error raises on the first attempt, never retried against the increment.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + from prisma.errors import UniqueViolationError + + non_deadlock = UniqueViolationError( + data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "LiteLLM_VerificationToken"}}} + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(non_deadlock)) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(UniqueViolationError): + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + mock_prisma_client.db.tx.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_spend_retries_deadlock(monkeypatch): + """The daily-spend upsert path retries a deadlock on the bulk upsert and then drains successfully.""" + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[_deadlock_error(), None]) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + daily_spend_transactions = {"k1": _daily_txn()} + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert mock_prisma_client.db.execute_raw.call_count == 2 + assert daily_spend_transactions == {} + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.parametrize( + "transactions_key, sample_key", + [ + ("user_list_transactions", "user-1"), + ("team_list_transactions", "team-1"), + ("team_member_list_transactions", "team_id::team-1::user_id::user-1"), + ("org_list_transactions", "org-1"), + ("tag_list_transactions", "tag-1"), + ("agent_list_transactions", "agent-1"), + ], +) +@pytest.mark.asyncio +async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkeypatch, transactions_key, sample_key): + """Every per-entity spend path, not just keys, retries a deadlock instead of dropping the increment.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + mock_batcher = MagicMock() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + proxy_logging.call_details = {} + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(**{transactions_key: {sample_key: 0.5}}), + ) + + assert mock_prisma_client.db.tx.call_count == 2 + proxy_logging.failure_handler.assert_not_called() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index a188289bfce..474e571e592 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -549,3 +549,36 @@ def test_handle_db_exception_surfaces_a_permanent_fault_even_when_degraded_mode_ with pytest.raises(BinaryNotFoundError): PrismaDBExceptionHandler.handle_db_exception(BinaryNotFoundError("query engine binary not found")) + + +@pytest.mark.parametrize( + "error", + [ + RawQueryError(data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "t"}}}), + PrismaError("Transaction failed due to a write conflict or a deadlock. Please retry your transaction"), + RawQueryError(data={"user_facing_error": {"message": "deadlock detected", "meta": {"table": "t"}}}), + RawQueryError( + data={"user_facing_error": {"message": "ERROR: 40P01: deadlock detected", "meta": {"table": "t"}}} + ), + ], +) +def test_is_deadlock_error_matches_postgres_deadlock(error): + """A Postgres deadlock surfaced through prisma (P2034 or 40P01 / "deadlock detected" text) is recognized.""" + assert PrismaDBExceptionHandler.is_deadlock_error(error) is True + + +@pytest.mark.parametrize( + "error", + [ + UniqueViolationError(data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + PrismaError("validation failed on query"), + PrismaError("can't reach database server"), + httpx.ConnectError("connection refused"), + RuntimeError("deadlock detected"), + ValueError("40P01"), + ], +) +def test_is_deadlock_error_excludes_non_deadlocks(error): + """Non-deadlock prisma errors, connectivity failures, and non-prisma exceptions are not treated as deadlocks.""" + assert PrismaDBExceptionHandler.is_deadlock_error(error) is False From 16e6ad6fb73566b1cbf3aa04fe90f7ee4653a389 Mon Sep 17 00:00:00 2001 From: RayJueWang <570828708@qq.com> Date: Thu, 6 Aug 2026 16:27:33 +0800 Subject: [PATCH 13/88] fix(proxy): recognize P2034 write-conflict deadlock text in is_deadlock_error The prisma P2034 transaction conflict can surface only as the message "Transaction failed due to a write conflict or a deadlock" without the code being reachable on the raised object, so the message fallback in is_deadlock_error now matches that canonical wording in addition to 40P01 / deadlock detected. Fixes the proxy-infra unit test that asserts this exact prisma message is treated as a retryable deadlock. --- litellm/proxy/db/exception_handler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 91c7e576dff..f7a39aaa50f 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -176,7 +176,11 @@ class PrismaDBExceptionHandler: if getattr(e, "code", None) == "P2034": return True error_message = str(e).lower() - return "deadlock detected" in error_message or "40p01" in error_message + return ( + "deadlock detected" in error_message + or "40p01" in error_message + or "write conflict or a deadlock" in error_message + ) @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: From 28c1e431968d61ea7caf1d82e351aca765ec83a1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 14 Aug 2026 00:01:16 -0700 Subject: [PATCH 14/88] feat(ui): standardize the Teams page header --- .../_components/AccessGroupsPage.tsx | 4 +- .../budgets/_components/budget_panel.tsx | 4 +- .../projects/_components/ProjectsPage.tsx | 4 +- .../src/components/Teams.test.tsx | 37 ++++++--- ui/litellm-dashboard/src/components/Teams.tsx | 47 +++++------ .../VirtualKeysPage/VirtualKeysTable.tsx | 4 +- .../shared/LegacyPageHeader.test.tsx | 33 ++++++++ .../components/shared/LegacyPageHeader.tsx | 25 ++++++ .../src/components/shared/PageHeader.test.tsx | 80 +++++++++++++++---- .../src/components/shared/PageHeader.tsx | 63 +++++++++++---- 10 files changed, 224 insertions(+), 77 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index f37acb3d85a..aeff249fdd3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -3,7 +3,7 @@ import { useDeleteAccessGroup } from "@/app/(dashboard)/hooks/accessGroups/useDe import { Plus, SearchIcon, X } from "lucide-react"; import { useMemo, useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; @@ -61,7 +61,7 @@ export function AccessGroupsPage() { return (
- = ({ accessToken }) => { return (
- } title="Budgets" subtitle="Spend, TPM and RPM limits you can assign to customers." diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx index 91ad9f847f5..dc18a05edca 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx @@ -3,7 +3,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Plus, SearchIcon, X } from "lucide-react"; import { parseAsString, useQueryState } from "nuqs"; import { useMemo, useState } from "react"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { CreateProjectModal } from "./ProjectModals/CreateProjectModal"; @@ -56,7 +56,7 @@ export function ProjectsPage() { return (
- { expect(onUrlUpdate.mock.calls.at(-1)![0].searchParams.has("team")).toBe(false); await waitFor(() => expect(screen.queryByTestId("team-info-view")).not.toBeInTheDocument()); }); + + it("should preserve the legacy inset for the team detail view", async () => { + renderWithQueryClient(, { + searchParams: "?team=team-from-url", + }); + + await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); + expect(screen.getByRole("main")).toHaveClass("px-12", "py-6"); + }); }); describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { @@ -521,22 +530,28 @@ describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { mockUseOrganizations.mockReturnValue({ data: [] }); }); - it("renders the Create Team button inside the tab bar, ahead of the tabs", () => { - const { container } = renderWithQueryClient(); + it("should render the Create Team button inside the tab bar, ahead of the tabs", () => { + renderWithQueryClient(); - const createButton = screen.getByTestId("create-team-button"); - const tabNav = container.querySelector(".ant-tabs-nav"); + const tabNav = screen.getByRole("tablist"); + const createButton = within(tabNav).getByTestId("create-team-button"); + const firstTab = within(tabNav).getByRole("tab", { name: "Your Teams" }); + const tabs = tabNav.closest(".ant-tabs"); - // The CTA lives in the tab bar's left slot, not the standalone page header. - expect(tabNav).not.toBeNull(); - expect(tabNav!.contains(createButton)).toBe(true); - - // It reads as the left end of the cluster: it precedes the first tab in DOM order. - const firstTab = screen.getByRole("tab", { name: "Your Teams" }); + expect(screen.getByRole("main")).toHaveClass("p-8"); + expect(within(tabNav).getByRole("separator")).toBeInTheDocument(); expect(createButton.compareDocumentPosition(firstTab) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(tabs).toHaveClass( + "[&>.ant-tabs-nav]:!mb-6", + "[&>.ant-tabs-nav]:before:!border-b-0", + "[&_.ant-tabs-ink-bar]:!h-0.5", + "[&_.ant-tabs-tab]:!py-[7px]", + "[&_.ant-tabs-tab+_.ant-tabs-tab]:!ml-[22px]", + "[&_.ant-tabs-tab-active]:font-semibold", + ); }); - it("omits the Create Team CTA for a role that cannot manage teams", () => { + it("should omit the Create Team CTA for a role that cannot manage teams", () => { renderWithQueryClient(); expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index becbe0e2b48..5b79b067412 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -6,7 +6,7 @@ import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/react"; -import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; +import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, Tooltip, Typography } from "antd"; import { Plus, Users } from "lucide-react"; import React, { useEffect, useState } from "react"; import { useQuery, useQueryClient } from "@tanstack/react-query"; @@ -403,7 +403,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser return false; }; - const { token } = theme.useToken(); const { Text } = Typography; const { Content } = Layout; @@ -474,7 +473,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ]; return ( - + {selectedTeamId ? ( = ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} /> ) : ( - <> -
- } - title="Teams" - subtitle="Manage teams, members, and their access to models and budgets" + } + title="Teams" + subtitle="Manage teams, members, and their access to models and budgets" + primaryAction={ + canCreateOrManageTeams(userRole, userID, organizations) ? ( + setIsTeamModalVisible(true)} data-testid="create-team-button"> + + Create Team + + ) : undefined + } + tabs={({ leadingControls }) => ( + -
- - - setIsTeamModalVisible(true)} data-testid="create-team-button"> - - Create Team - -
-
- ) : undefined, - }} - /> - + )} + /> )} {canCreateOrManageTeams(userRole, userID, organizations) && ( diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index fa0360c0dda..b278b4f675d 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -12,7 +12,7 @@ import { DataTableToolbar, } from "@/components/shared/DataTable"; import { SearchSelect } from "@/components/shared/SearchSelect"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Input } from "@/components/ui/input"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; @@ -172,7 +172,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { return (
- } title="Virtual Keys" subtitle="Every key that authenticates requests to the gateway." diff --git a/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx new file mode 100644 index 00000000000..a0081c1c7f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx @@ -0,0 +1,33 @@ +import { renderWithProviders, screen } from "@/../tests/test-utils"; +import { describe, expect, it } from "vitest"; + +import { LegacyPageHeader } from "./LegacyPageHeader"; + +describe("LegacyPageHeader", () => { + it("should render the title as a heading", () => { + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + + it("should render the optional identity and actions", () => { + renderWithProviders( + Key icon} + actions={} + />, + ); + + expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); + expect(screen.getByText("Key icon")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + }); + + it("should omit optional actions when none are provided", () => { + renderWithProviders(); + + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx new file mode 100644 index 00000000000..43979ad00b4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx @@ -0,0 +1,25 @@ +"use client"; + +import * as React from "react"; + +interface LegacyPageHeaderProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + icon?: React.ReactNode; + actions?: React.ReactNode; +} + +export function LegacyPageHeader({ title, subtitle, icon, actions }: LegacyPageHeaderProps) { + return ( +
+
+ {icon != null && {icon}} +
+

{title}

+ {subtitle != null &&

{subtitle}

} +
+
+ {actions != null &&
{actions}
} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx index f7a313271da..3741542abad 100644 --- a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx @@ -1,31 +1,77 @@ -import { render, screen } from "@testing-library/react"; +import { renderWithProviders, screen, within } from "@/../tests/test-utils"; import { describe, expect, it } from "vitest"; import { PageHeader } from "./PageHeader"; +const identity = { + icon: Teams icon, + title: "Teams", + subtitle: "Manage teams, members, and their access to models and budgets", +}; + describe("PageHeader", () => { - it("renders the title as a heading", () => { - render(); - expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + it("should render the page identity", () => { + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "Teams" })).toBeInTheDocument(); + expect(screen.getByText("Teams icon").parentElement).toHaveAttribute("aria-hidden", "true"); + expect(screen.getByText(identity.subtitle)).toBeInTheDocument(); }); - it("renders the subtitle, icon, and actions when provided", () => { - render( + it("should apply the standard title and subtext typography", () => { + renderWithProviders(); + + const icon = screen.getByText("Teams icon").parentElement; + expect(screen.getByRole("heading", { name: "Teams" })).toHaveClass("text-2xl", "font-semibold", "tracking-tight"); + expect(screen.getByText(identity.subtitle)).toHaveClass("mt-1.5", "text-sm", "text-muted-foreground"); + expect(icon).toHaveClass("size-5", "[&_svg]:size-5", "[&_svg]:stroke-[1.75]"); + expect(icon?.parentElement).toHaveClass("gap-2.5"); + }); + + it("should render the primary action, divider, tabs, and utilities in the standard control row", () => { + renderWithProviders( } - actions={} + {...identity} + primaryAction={} + tabs={ +
+ +
+ } + utilities={} />, ); - expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); - expect(screen.getByTestId("icon")).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + + const controls = screen.getByRole("group", { name: "Page controls" }); + expect(controls).toHaveClass("mt-5", "h-9"); + expect(within(controls).getByRole("separator")).toHaveClass("mx-4", "h-6"); + expect(controls).toHaveTextContent("Create TeamYour TeamsRefresh"); }); - it("omits the optional slots when not provided", () => { - render(); - expect(screen.queryByRole("button")).not.toBeInTheDocument(); - expect(document.querySelector("p")).toBeNull(); + it("should omit the divider when tabs are absent", () => { + renderWithProviders(Create Team} />); + + expect(screen.queryByRole("separator")).not.toBeInTheDocument(); + }); + + it("should provide standard controls to an embedded tab shell", () => { + renderWithProviders( + Create Team} + tabs={({ leadingControls, utilities }) => ( +
+ {leadingControls} + + {utilities} +
+ )} + utilities={} + />, + ); + + const tabs = screen.getByRole("tablist"); + expect(within(tabs).getByRole("separator")).toBeInTheDocument(); + expect(tabs).toHaveTextContent("Create TeamYour TeamsRefresh"); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx index e314e8e8bc2..81092821efc 100644 --- a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx @@ -2,24 +2,57 @@ import * as React from "react"; -interface PageHeaderProps { - title: React.ReactNode; - subtitle?: React.ReactNode; - icon?: React.ReactNode; - actions?: React.ReactNode; +import { ToolbarSeparator } from "./ToolbarSeparator"; + +interface EmbeddedTabsSlots { + leadingControls: React.ReactNode; + utilities: React.ReactNode; } -export function PageHeader({ title, subtitle, icon, actions }: PageHeaderProps) { - return ( -
-
- {icon != null && {icon}} -
-

{title}

- {subtitle != null &&

{subtitle}

} -
+interface PageHeaderProps { + title: React.ReactNode; + subtitle: React.ReactNode; + icon: React.ReactNode; + primaryAction?: React.ReactNode; + tabs?: React.ReactNode | ((slots: EmbeddedTabsSlots) => React.ReactNode); + utilities?: React.ReactNode; +} + +export function PageHeader({ title, subtitle, icon, primaryAction, tabs, utilities }: PageHeaderProps) { + const leadingControls = + primaryAction == null ? null : ( +
+ {primaryAction} + {tabs != null && }
- {actions != null &&
{actions}
} + ); + const utilityControls = utilities == null ? null :
{utilities}
; + const hasControlRow = primaryAction != null || tabs != null || utilities != null; + + return ( +
+
+ +

{title}

+
+

{subtitle}

+ + {typeof tabs === "function" ? ( +
{tabs({ leadingControls, utilities: utilityControls })}
+ ) : ( + hasControlRow && ( +
+ {leadingControls} + {tabs} + {utilityControls != null &&
{utilityControls}
} +
+ ) + )}
); } From ed33687422c544afa7ba6268294744b478026557 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:11:05 -0700 Subject: [PATCH 15/88] feat(proxy): auto-suppress the no-Redis banner for confirmed single-worker deployments --- .../migration.sql | 9 ++ .../litellm_proxy_extras/schema.prisma | 11 ++ litellm/proxy/db/proxy_worker_heartbeat.py | 89 +++++++++++++++ .../health_endpoints/_health_endpoints.py | 19 +++- litellm/proxy/proxy_server.py | 32 +++++- litellm/proxy/schema.prisma | 11 ++ schema.prisma | 11 ++ .../proxy/db/test_proxy_worker_heartbeat.py | 81 ++++++++++++++ .../health_endpoints/test_health_endpoints.py | 105 +++++++++++++++--- .../components/NoRedisWarningBanner.test.tsx | 1 + .../src/components/NoRedisWarningBanner.tsx | 8 +- 11 files changed, 349 insertions(+), 28 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql create mode 100644 litellm/proxy/db/proxy_worker_heartbeat.py create mode 100644 tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql new file mode 100644 index 00000000000..0a5d9df8aaf --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql @@ -0,0 +1,9 @@ +-- CreateTable +CREATE TABLE "LiteLLM_ProxyWorkerHeartbeat" ( + "worker_id" TEXT NOT NULL, + "hostname" TEXT NOT NULL, + "started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "last_heartbeat_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_ProxyWorkerHeartbeat_pkey" PRIMARY KEY ("worker_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py new file mode 100644 index 00000000000..6a2a4572e43 --- /dev/null +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -0,0 +1,89 @@ +""" +Live proxy worker census, one row per worker process. + +Every uvicorn worker upserts its own row on a fixed heartbeat, so counting +rows with a recent heartbeat answers "how many workers share this database?" +without any coordination. The Admin UI's "no Redis" banner uses that count to +hide itself for deployments that are provably a single worker, where per-worker +rate limits, budgets, and router state are already global. All timestamps are +written and compared with the database's own clock, so pods with skewed clocks +still agree. +""" + +from __future__ import annotations + +import socket +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS: Final = 60 +PROXY_WORKER_LIVENESS_WINDOW_SECONDS: Final = 3 * PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS +STALE_ROW_RETENTION_SECONDS: Final = 3600 + +BEAT_SQL: Final = """ +INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at) +VALUES ($1, $2, NOW()) +ON CONFLICT (worker_id) DO UPDATE SET last_heartbeat_at = NOW() +""" + +PRUNE_SQL: Final = """ +DELETE FROM "LiteLLM_ProxyWorkerHeartbeat" +WHERE last_heartbeat_at < NOW() - make_interval(secs => $1) +""" + +COUNT_SQL: Final = """ +SELECT COUNT(*)::int AS live_workers FROM "LiteLLM_ProxyWorkerHeartbeat" +WHERE last_heartbeat_at > NOW() - make_interval(secs => $1) +""" + +DEREGISTER_SQL: Final = """ +DELETE FROM "LiteLLM_ProxyWorkerHeartbeat" WHERE worker_id = $1 +""" + + +class _LiveWorkerCountRow(TypedDict): + live_workers: ReadOnly[int] + + +_COUNT_ROWS_ADAPTER: Final = TypeAdapter(tuple[_LiveWorkerCountRow, ...]) + + +class ProxyWorkerHeartbeat: + def __init__(self, prisma_client: PrismaClient, worker_id: str | None = None) -> None: + self.prisma_client: Final = prisma_client + self.worker_id: Final[str] = worker_id or str(uuid.uuid4()) + self.hostname: Final = socket.gethostname() + + async def beat(self) -> None: + try: + await self.prisma_client.db.execute_raw(BEAT_SQL, self.worker_id, self.hostname) + await self.prisma_client.db.execute_raw(PRUNE_SQL, STALE_ROW_RETENTION_SECONDS) + except Exception as beat_err: # noqa: BLE001 # a missed heartbeat must never take down the worker + verbose_proxy_logger.debug("Proxy worker heartbeat write failed: %s", beat_err) + + async def deregister(self) -> None: + try: + await self.prisma_client.db.execute_raw(DEREGISTER_SQL, self.worker_id) + except Exception as deregister_err: # noqa: BLE001 # best-effort cleanup; the liveness window ages the row out anyway + verbose_proxy_logger.debug("Proxy worker heartbeat deregister failed: %s", deregister_err) + + +async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: + """ + The number of workers with a recent heartbeat, or None when the database + cannot answer. Callers must treat None as "unknown", not as zero. + """ + try: + rows: Final = await prisma_client.db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] + except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 + verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) + return None diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e814ec42d26..33894777bc3 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, @@ -1451,7 +1452,7 @@ def callback_name(callback): DISABLE_NO_REDIS_WARNING_ENV_VAR: Final = "LITELLM_DISABLE_NO_REDIS_WARNING" -def _show_no_redis_warning() -> bool: +async def _show_no_redis_warning() -> bool: """ Whether the UI should warn that no Redis is configured. @@ -1461,16 +1462,22 @@ def _show_no_redis_warning() -> bool: coordination cache (from a Redis response cache, general_settings. coordination_redis, or the REDIS_* env fallback) and the router's own Redis (router_settings.redis_host), which backs cooldowns and usage-based - routing on its own. Operators who know they run one worker can silence the - warning with LITELLM_DISABLE_NO_REDIS_WARNING=true. + routing on its own. A deployment whose worker-heartbeat census proves it + is exactly one worker needs no cross-worker coordination, so it never + warns; when the census is unavailable or shows more than one worker, the + warning stands unless LITELLM_DISABLE_NO_REDIS_WARNING=true silences it. """ - from litellm.proxy.proxy_server import llm_router, redis_usage_cache + from litellm.proxy.proxy_server import llm_router, prisma_client, redis_usage_cache if redis_usage_cache is not None: return False if llm_router is not None and llm_router.cache.redis_cache is not None: return False - return get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is not True + if get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is True: + return False + if prisma_client is None: + return True + return await count_live_proxy_workers(prisma_client) != 1 async def _get_health_readiness_details( @@ -1513,7 +1520,7 @@ async def _get_health_readiness_details( # check log level log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) - show_no_redis_warning: Final = _show_no_redis_warning() + show_no_redis_warning: Final = await _show_no_redis_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 359187f81cb..6ee08a732f2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -379,6 +379,10 @@ from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, flush_gateway_requests, ) +from litellm.proxy.db.proxy_worker_heartbeat import ( + PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, + ProxyWorkerHeartbeat, +) from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router @@ -864,9 +868,11 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e) -async def proxy_shutdown_event(): +async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None): global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") + if worker_heartbeat is not None and prisma_client: + await worker_heartbeat.deregister() if prisma_client: # Drain the SGR fold first: it lives in memory, so an un-drained interval # is lost, and a write attempted after disconnect raises @@ -1200,7 +1206,7 @@ async def proxy_startup_event(app: FastAPI): ) ### START BATCH WRITING DB + CHECKING NEW MODELS### - if prisma_client is not None: + worker_heartbeat: Final = ( await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings=general_settings, prisma_client=prisma_client, @@ -1209,7 +1215,10 @@ async def proxy_startup_event(app: FastAPI): proxy_batch_write_at=proxy_batch_write_at, proxy_logging_obj=proxy_logging_obj, ) - + if prisma_client is not None + else None + ) + if prisma_client is not None: await ProxyStartupEvent._update_default_team_member_budget() ## SYNC UI SETTINGS ## @@ -1280,7 +1289,7 @@ async def proxy_startup_event(app: FastAPI): await proxy_config.stop_auth_cache_invalidation_subscriber() - await proxy_shutdown_event() + await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) def _generate_stable_operation_id(route: "APIRoute") -> str: @@ -8665,7 +8674,7 @@ class ProxyStartupEvent: proxy_budget_rescheduler_max_time: int, proxy_batch_write_at: int, proxy_logging_obj: ProxyLogging, - ): + ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" global store_model_in_db, scheduler @@ -8710,6 +8719,18 @@ class ProxyStartupEvent: # Ensure minimum interval of 30 seconds for batch writing to prevent memory issues batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5) + ### PROXY WORKER HEARTBEAT ### + worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client) + await worker_heartbeat.beat() + scheduler.add_job( + worker_heartbeat.beat, + "interval", + seconds=PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, + id="proxy_worker_heartbeat_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### RESET BUDGET ### if general_settings.get("disable_reset_budget", False) is False: budget_reset_job: Final = ResetBudgetJob( @@ -9048,6 +9069,7 @@ class ProxyStartupEvent: "APScheduler started with memory leak prevention settings: removed jitter, increased intervals, misfire_grace_time=%s", APSCHEDULER_MISFIRE_GRACE_TIME, ) + return worker_heartbeat @classmethod async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/schema.prisma b/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py new file mode 100644 index 00000000000..2209be0dc2e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py @@ -0,0 +1,81 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.proxy_worker_heartbeat import ( + BEAT_SQL, + COUNT_SQL, + DEREGISTER_SQL, + PROXY_WORKER_LIVENESS_WINDOW_SECONDS, + PRUNE_SQL, + STALE_ROW_RETENTION_SECONDS, + ProxyWorkerHeartbeat, + count_live_proxy_workers, +) + + +def _prisma(): + prisma = MagicMock() + prisma.db.execute_raw = AsyncMock() + prisma.db.query_raw = AsyncMock() + return prisma + + +@pytest.mark.asyncio +async def test_beat_upserts_own_row_then_prunes_stale_rows(): + prisma = _prisma() + heartbeat = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1") + await heartbeat.beat() + calls = prisma.db.execute_raw.call_args_list + assert calls[0].args == (BEAT_SQL, "worker-1", heartbeat.hostname) + assert calls[1].args == (PRUNE_SQL, STALE_ROW_RETENTION_SECONDS) + + +@pytest.mark.asyncio +async def test_beat_survives_a_database_error(): + prisma = _prisma() + prisma.db.execute_raw = AsyncMock(side_effect=RuntimeError("db down")) + await ProxyWorkerHeartbeat(prisma_client=prisma).beat() + + +def test_each_worker_process_gets_its_own_id(): + prisma = _prisma() + first = ProxyWorkerHeartbeat(prisma_client=prisma) + second = ProxyWorkerHeartbeat(prisma_client=prisma) + assert first.worker_id != second.worker_id + + +@pytest.mark.asyncio +async def test_deregister_deletes_only_its_own_row(): + prisma = _prisma() + await ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1").deregister() + assert prisma.db.execute_raw.call_args.args == (DEREGISTER_SQL, "worker-1") + + +@pytest.mark.asyncio +async def test_deregister_survives_a_database_error(): + prisma = _prisma() + prisma.db.execute_raw = AsyncMock(side_effect=RuntimeError("db down")) + await ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1").deregister() + + +@pytest.mark.asyncio +async def test_count_reads_workers_within_the_liveness_window(): + prisma = _prisma() + prisma.db.query_raw.return_value = [{"live_workers": 3}] + assert await count_live_proxy_workers(prisma) == 3 + assert prisma.db.query_raw.call_args.args == (COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + + +@pytest.mark.asyncio +async def test_count_returns_unknown_when_the_query_fails(): + prisma = _prisma() + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await count_live_proxy_workers(prisma) is None + + +@pytest.mark.asyncio +async def test_count_returns_unknown_for_a_malformed_row(): + prisma = _prisma() + prisma.db.query_raw.return_value = [{"unexpected": "shape"}] + assert await count_live_proxy_workers(prisma) is None diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e2705bd5fec..831f659051c 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2467,61 +2467,140 @@ class TestNoRedisWarning: def _router(redis_cache): return SimpleNamespace(cache=SimpleNamespace(redis_cache=redis_cache)) - def test_warns_when_no_redis_is_configured(self, monkeypatch): + @staticmethod + def _prisma_with_workers(live_workers=None, error=None): + prisma = MagicMock() + if error is not None: + prisma.db.query_raw = AsyncMock(side_effect=error) + else: + prisma.db.query_raw = AsyncMock(return_value=[{"live_workers": live_workers}]) + return prisma + + @pytest.mark.asyncio + async def test_warns_when_no_redis_and_no_db_to_count_workers(self, monkeypatch): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", None), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True - def test_warns_when_there_is_no_router_at_all(self, monkeypatch): + @pytest.mark.asyncio + async def test_warns_when_there_is_no_router_at_all(self, monkeypatch): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True - def test_stays_quiet_when_a_coordination_redis_is_configured(self, monkeypatch): + @pytest.mark.asyncio + async def test_stays_quiet_for_a_confirmed_single_worker(self, monkeypatch): + """One live worker needs no cross-worker coordination, so no env var is needed.""" monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(1)), + ): + assert await _show_no_redis_warning() is False + + @pytest.mark.asyncio + @pytest.mark.parametrize("live_workers", [2, 5]) + async def test_warns_when_multiple_workers_share_the_db(self, monkeypatch, live_workers): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(live_workers)), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_warns_when_the_worker_census_is_empty(self, monkeypatch): + """Zero rows means the census cannot CONFIRM a single worker, so warn.""" + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(0)), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_warns_when_the_worker_census_query_fails(self, monkeypatch): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch( + "litellm.proxy.proxy_server.prisma_client", + self._prisma_with_workers(error=RuntimeError("db down")), + ), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_stays_quiet_when_a_coordination_redis_is_configured(self, monkeypatch): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + prisma = self._prisma_with_workers(5) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", MagicMock()), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", prisma), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False + prisma.db.query_raw.assert_not_called() - def test_stays_quiet_when_only_the_router_has_redis(self, monkeypatch): + @pytest.mark.asyncio + async def test_stays_quiet_when_only_the_router_has_redis(self, monkeypatch): """router_settings.redis_host alone backs cooldowns and usage-based routing.""" monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(MagicMock())), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(5)), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False + @pytest.mark.asyncio @pytest.mark.parametrize("value", ["true", "True"]) - def test_env_var_suppresses_the_warning(self, monkeypatch, value): + async def test_env_var_suppresses_the_warning_despite_multiple_workers(self, monkeypatch, value): monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", value) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(5)), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False - def test_env_var_set_false_keeps_the_warning(self, monkeypatch): + @pytest.mark.asyncio + async def test_env_var_set_false_keeps_the_warning_for_multiple_workers(self, monkeypatch): monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", "false") with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(2)), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_env_var_set_false_does_not_force_the_warning_for_a_single_worker(self, monkeypatch): + monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", "false") + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(1)), + ): + assert await _show_no_redis_warning() is False @pytest.mark.asyncio @pytest.mark.parametrize("has_prisma_client", [True, False]) async def test_readiness_details_carries_the_flag(self, monkeypatch, has_prisma_client): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) - prisma_client = MagicMock() if has_prisma_client else None + prisma_client = self._prisma_with_workers(2) if has_prisma_client else None with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.redis_usage_cache", None), diff --git a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx index 8afde8eec94..600315e789e 100644 --- a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx @@ -20,6 +20,7 @@ describe("NoRedisWarningBanner", () => { renderWithProviders(); expect(screen.getByRole("alert")).toBeInTheDocument(); expect(screen.getByText(/No Redis configured\. Redis is highly recommended/i)).toBeInTheDocument(); + expect(screen.getByText(/more than one worker/i)).toBeInTheDocument(); }); it("should link to the docs page listing what breaks without Redis", () => { diff --git a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx index 93c0f55486d..02433fad521 100644 --- a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx +++ b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx @@ -26,13 +26,13 @@ export const NoRedisWarningBanner: React.FC = ({ acce

No Redis configured. Redis is highly recommended

- Rate limits, budgets, router state, and cache invalidation are per worker without Redis, so limits are - enforced once per worker and spend can overshoot.{" "} + This proxy is running more than one worker (or the worker count could not be verified). Without Redis, rate + limits, budgets, router state, and cache invalidation are per worker, so limits are enforced once per worker + and spend can overshoot.{" "} See everything that does not work without Redis - . If you run a single worker and this is intentional, set{" "} - LITELLM_DISABLE_NO_REDIS_WARNING=true to hide this banner. + . Set LITELLM_DISABLE_NO_REDIS_WARNING=true to hide this banner anyway.

From 3217b8edae27074298717b2a56fd1d5a82b6d517 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:23:06 -0700 Subject: [PATCH 16/88] fix(proxy): count worker heartbeats on the primary so replica lag cannot undercount --- litellm/proxy/db/proxy_worker_heartbeat.py | 8 ++++++-- .../proxy/db/test_proxy_worker_heartbeat.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 6a2a4572e43..990ff48eb18 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -20,6 +20,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -79,10 +80,13 @@ class ProxyWorkerHeartbeat: async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: """ The number of workers with a recent heartbeat, or None when the database - cannot answer. Callers must treat None as "unknown", not as zero. + cannot answer. Callers must treat None as "unknown", not as zero. Always + counts on the primary: a lagging read replica must never undercount. """ try: - rows: Final = await prisma_client.db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + db: Final = prisma_client.db + primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db + rows: Final = await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) diff --git a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py index 2209be0dc2e..33ae6190411 100644 --- a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py +++ b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py @@ -12,6 +12,7 @@ from litellm.proxy.db.proxy_worker_heartbeat import ( ProxyWorkerHeartbeat, count_live_proxy_workers, ) +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper def _prisma(): @@ -67,6 +68,18 @@ async def test_count_reads_workers_within_the_liveness_window(): assert prisma.db.query_raw.call_args.args == (COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) +@pytest.mark.asyncio +async def test_count_reads_from_the_primary_when_reads_route_to_a_replica(): + writer = MagicMock() + writer.query_raw = AsyncMock(return_value=[{"live_workers": 2}]) + reader = MagicMock() + reader.query_raw = AsyncMock(return_value=[{"live_workers": 1}]) + prisma = MagicMock() + prisma.db = RoutingPrismaWrapper(writer=writer, reader=reader) + assert await count_live_proxy_workers(prisma) == 2 + reader.query_raw.assert_not_awaited() + + @pytest.mark.asyncio async def test_count_returns_unknown_when_the_query_fails(): prisma = _prisma() From 3d523d6d81816d692927e00a44ad998665731b82 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:14:42 +0000 Subject: [PATCH 17/88] fix(model_prices): add provider-announced deprecation_date to 205 registry entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 205 ++++++++++++++++++ model_prices_and_context_window.json | 205 ++++++++++++++++++ 2 files changed, 410 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78b53cefc53..1a130c4ac0a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54,6 +54,7 @@ "output_cost_per_image": 0.04 }, "1024-x-1024/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 1.9e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -67,6 +68,7 @@ "output_cost_per_image": 0.08 }, "256-x-256/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 2.4414e-07, "litellm_provider": "openai", "mode": "image_generation", @@ -80,6 +82,7 @@ "output_cost_per_image": 0.018 }, "512-x-512/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.86e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -2887,6 +2890,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -2908,6 +2912,7 @@ "supports_vision": true }, "azure_ai/claude-opus-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -2930,6 +2935,7 @@ "supports_output_config": true }, "azure_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -2959,6 +2965,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-06", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -3083,6 +3090,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -3104,6 +3112,7 @@ "supports_vision": true }, "azure_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3156,6 +3165,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-sonnet-4-6": { + "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -3226,6 +3236,7 @@ "supports_tool_choice": true }, "azure_ai/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -3318,6 +3329,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3364,6 +3376,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-2026-03-05": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3410,6 +3423,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3455,6 +3469,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro-2026-03-05": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3500,6 +3515,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3540,6 +3556,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-mini-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3580,6 +3597,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3620,6 +3638,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3849,6 +3868,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3918,6 +3938,7 @@ "supports_none_reasoning_effort": true }, "azure/eu/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3948,6 +3969,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -4107,6 +4129,7 @@ "supports_vision": true }, "azure/global-standard/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4155,6 +4178,7 @@ "supports_vision": true }, "azure/global/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4224,6 +4248,7 @@ "supports_none_reasoning_effort": true }, "azure/global/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4254,6 +4279,7 @@ "supports_vision": true }, "azure/global/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -4492,6 +4518,7 @@ "supports_vision": true }, "azure/gpt-4.1": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4559,6 +4586,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4626,6 +4654,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4902,6 +4931,7 @@ "supports_vision": false }, "azure/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", @@ -5344,6 +5374,7 @@ "supports_vision": true }, "azure/gpt-5": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5507,6 +5538,7 @@ "supports_vision": true }, "azure/gpt-5-mini": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5572,6 +5604,7 @@ "supports_vision": true }, "azure/gpt-5-nano": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 5e-09, "input_cost_per_token": 5e-08, "litellm_provider": "azure", @@ -5667,6 +5700,7 @@ "supports_vision": true }, "azure/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5736,6 +5770,7 @@ "supports_none_reasoning_effort": true }, "azure/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5797,6 +5832,7 @@ "supports_vision": true }, "azure/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5827,6 +5863,7 @@ "supports_vision": true }, "azure/gpt-5.2": { + "deprecation_date": "2027-06-08", "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", @@ -6136,6 +6173,7 @@ "supports_web_search": true }, "azure/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6180,6 +6218,7 @@ "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6218,6 +6257,7 @@ "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6379,6 +6419,7 @@ "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -7045,6 +7086,7 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -7095,6 +7137,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7142,6 +7185,7 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7408,6 +7452,7 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7489,6 +7534,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7601,6 +7647,7 @@ "output_cost_per_token": 0.0 }, "azure/high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7610,6 +7657,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7619,6 +7667,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7628,6 +7677,7 @@ ] }, "azure/low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7637,6 +7687,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7646,6 +7697,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7655,6 +7707,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7664,6 +7717,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7673,6 +7727,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7695,6 +7750,7 @@ ] }, "azure/gpt-image-1.5": { + "deprecation_date": "2027-06-16", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7720,6 +7776,7 @@ ] }, "azure/gpt-image-2": { + "deprecation_date": "2027-10-21", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7751,6 +7808,7 @@ "supports_pdf_input": true }, "azure/low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7760,6 +7818,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7769,6 +7828,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0345052083e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7778,6 +7838,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7787,6 +7848,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7796,6 +7858,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 7.9752604167e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7805,6 +7868,7 @@ ] }, "azure/high/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7814,6 +7878,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7823,6 +7888,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.1575520833e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7850,6 +7916,7 @@ "supports_function_calling": true }, "azure/o1": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -7944,6 +8011,7 @@ "supports_vision": false }, "azure/o3": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8041,6 +8109,7 @@ "supports_web_search": true }, "azure/o3-mini": { + "deprecation_date": "2026-10-01", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8071,6 +8140,7 @@ "supports_vision": false }, "azure/o3-pro": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8132,6 +8202,7 @@ "supports_vision": true }, "azure/o4-mini": { + "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8580,6 +8651,7 @@ "supports_vision": true }, "azure/us/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8649,6 +8721,7 @@ "supports_none_reasoning_effort": true }, "azure/us/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8679,6 +8752,7 @@ "supports_vision": true }, "azure/us/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -8876,6 +8950,7 @@ ] }, "azure_ai/FW-DeepSeek-V3.2": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, "input_cost_per_token": 6.2e-07, "litellm_provider": "azure_ai", @@ -8906,6 +8981,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure_ai", @@ -8921,6 +8997,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5.1": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.86e-07, "input_cost_per_token": 1.54e-06, "litellm_provider": "azure_ai", @@ -8987,6 +9064,7 @@ "supports_tool_choice": true }, "azure_ai/FW-Kimi-K2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 6.6e-07, "litellm_provider": "azure_ai", @@ -9079,6 +9157,7 @@ "supports_vision": true }, "azure_ai/FW-MiniMax-M2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, "input_cost_per_token": 3.3e-07, "litellm_provider": "azure_ai", @@ -9164,6 +9243,7 @@ ] }, "azure_ai/MAI-Image-2e": { + "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", "mode": "image_generation", @@ -9175,6 +9255,7 @@ ] }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9188,6 +9269,7 @@ "supports_vision": true }, "azure_ai/Llama-3.2-90B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 2.04e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9249,6 +9331,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-405B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 5.33e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9271,6 +9354,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9452,6 +9536,7 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2505": { + "deprecation_date": "2026-07-20", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", @@ -9529,6 +9614,7 @@ "output_cost_per_token": 0.0 }, "azure_ai/cohere-rerank-v3.5": { + "deprecation_date": "2026-05-14", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -9591,6 +9677,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-r1": { + "deprecation_date": "2026-08-13", "input_cost_per_token": 1.35e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9614,6 +9701,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3-0324": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.14e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9626,6 +9714,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.1": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.23e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9639,6 +9728,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.74e-06, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9652,6 +9742,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-flash": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.9e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9683,6 +9774,7 @@ "supports_embedding_image_input": true }, "azure_ai/global/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9697,6 +9789,7 @@ "supports_web_search": true }, "azure_ai/global/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9712,6 +9805,7 @@ "supports_web_search": true }, "azure_ai/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9726,6 +9820,7 @@ "supports_web_search": true }, "azure_ai/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9773,6 +9868,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-non-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9786,6 +9882,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9863,6 +9960,7 @@ "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { + "deprecation_date": "2027-01-26", "input_cost_per_token": 6e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -9877,6 +9975,7 @@ "supports_vision": true }, "azure_ai/kimi-k2.6": { + "deprecation_date": "2027-04-16", "input_cost_per_token": 9.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -10004,6 +10103,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -11999,6 +12099,7 @@ ] }, "claude-haiku-4-5-20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12022,6 +12123,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12170,6 +12272,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12203,6 +12306,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5-20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12237,6 +12341,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-5": { + "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -12273,6 +12378,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-6": { + "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12419,6 +12525,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12448,6 +12555,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12477,6 +12585,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12513,6 +12622,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6-20260205": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12549,6 +12659,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12587,6 +12698,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-opus-4-7-20260416": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12625,6 +12737,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { + "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -12660,6 +12773,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-5": { + "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12698,6 +12812,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-4-8": { + "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14801,6 +14916,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -18353,6 +18469,7 @@ } }, "gemini-2.5-flash": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18398,6 +18515,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18442,6 +18560,7 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -18522,6 +18641,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -18646,6 +18766,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -18702,6 +18823,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -18791,6 +18913,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, @@ -19062,6 +19185,7 @@ "supports_image_size": false }, "gemini-2.5-pro": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19373,6 +19497,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19809,6 +19934,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-robotics-er-1.6-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_audio_token": 2e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -19879,6 +20005,7 @@ "supports_vision": true }, "gemini-embedding-001": { + "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, @@ -21492,6 +21619,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -23004,6 +23132,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -24135,6 +24264,7 @@ "supports_pdf_input": true }, "low/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24146,6 +24276,7 @@ "supports_pdf_input": true }, "low/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24157,6 +24288,7 @@ "supports_pdf_input": true }, "low/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24168,6 +24300,7 @@ "supports_pdf_input": true }, "medium/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.034, "litellm_provider": "openai", "mode": "image_generation", @@ -24179,6 +24312,7 @@ "supports_pdf_input": true }, "medium/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24190,6 +24324,7 @@ "supports_pdf_input": true }, "medium/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24201,6 +24336,7 @@ "supports_pdf_input": true }, "high/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.133, "litellm_provider": "openai", "mode": "image_generation", @@ -24212,6 +24348,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24223,6 +24360,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24234,6 +24372,7 @@ "supports_pdf_input": true }, "standard/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24245,6 +24384,7 @@ "supports_pdf_input": true }, "standard/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24256,6 +24396,7 @@ "supports_pdf_input": true }, "standard/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24267,6 +24408,7 @@ "supports_pdf_input": true }, "1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24278,6 +24420,7 @@ "supports_pdf_input": true }, "1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24289,6 +24432,7 @@ "supports_pdf_input": true }, "1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -27202,18 +27346,21 @@ "output_cost_per_second": 0.0 }, "hd/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 7.629e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -27260,6 +27407,7 @@ "max_output_tokens": 8192 }, "high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.167, "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "openai", @@ -27270,6 +27418,7 @@ ] }, "high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -27280,6 +27429,7 @@ ] }, "high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -28067,6 +28217,7 @@ "supports_tool_choice": true }, "low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.011, "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "openai", @@ -28077,6 +28228,7 @@ ] }, "low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28087,6 +28239,7 @@ ] }, "low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28111,6 +28264,7 @@ "output_cost_per_image": 0.072 }, "medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.042, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28121,6 +28275,7 @@ ] }, "medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28131,6 +28286,7 @@ ] }, "medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28141,6 +28297,7 @@ ] }, "low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.005, "litellm_provider": "openai", "mode": "image_generation", @@ -28149,6 +28306,7 @@ ] }, "low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28157,6 +28315,7 @@ ] }, "low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28165,6 +28324,7 @@ ] }, "medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.011, "litellm_provider": "openai", "mode": "image_generation", @@ -28173,6 +28333,7 @@ ] }, "medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -28181,6 +28342,7 @@ ] }, "medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -30074,6 +30236,7 @@ ] }, "multimodalembedding@001": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2e-07, "input_cost_per_image": 0.0001, "input_cost_per_token": 8e-07, @@ -35772,18 +35935,21 @@ "output_cost_per_image": 0.14 }, "standard/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 3.81469e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -35847,6 +36013,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "text-embedding-005": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -35920,6 +36087,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35929,6 +36097,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35938,6 +36107,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35947,6 +36117,7 @@ "output_cost_per_token": 0.0 }, "text-multilingual-embedding-002": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -38434,6 +38605,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38457,6 +38629,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-haiku-4-5@20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38609,6 +38782,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38636,6 +38810,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38654,6 +38829,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-1@20250805": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38672,6 +38848,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38700,6 +38877,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-5@20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38729,6 +38907,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38759,6 +38938,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38789,6 +38969,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38820,6 +39001,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38851,6 +39033,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38882,6 +39065,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38913,6 +39097,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38945,6 +39130,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38977,6 +39163,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39009,6 +39196,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39041,6 +39229,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39069,6 +39258,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39131,6 +39321,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5@20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39160,6 +39351,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -39187,6 +39379,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39218,6 +39411,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39382,6 +39576,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -39427,6 +39622,7 @@ "supports_image_size": false }, "vertex_ai/gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -39459,6 +39655,7 @@ "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -39535,6 +39732,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -39591,6 +39789,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -40308,6 +40507,7 @@ "supports_tool_choice": true }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40322,6 +40522,7 @@ ] }, "vertex_ai/veo-3.0-fast-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40336,6 +40537,7 @@ ] }, "vertex_ai/veo-3.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40378,6 +40580,7 @@ ] }, "vertex_ai/veo-3.1-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40392,6 +40595,7 @@ ] }, "vertex_ai/veo-3.1-fast-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -46773,6 +46977,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78b53cefc53..1a130c4ac0a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54,6 +54,7 @@ "output_cost_per_image": 0.04 }, "1024-x-1024/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 1.9e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -67,6 +68,7 @@ "output_cost_per_image": 0.08 }, "256-x-256/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 2.4414e-07, "litellm_provider": "openai", "mode": "image_generation", @@ -80,6 +82,7 @@ "output_cost_per_image": 0.018 }, "512-x-512/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.86e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -2887,6 +2890,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -2908,6 +2912,7 @@ "supports_vision": true }, "azure_ai/claude-opus-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -2930,6 +2935,7 @@ "supports_output_config": true }, "azure_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -2959,6 +2965,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-06", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -3083,6 +3090,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -3104,6 +3112,7 @@ "supports_vision": true }, "azure_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3156,6 +3165,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-sonnet-4-6": { + "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -3226,6 +3236,7 @@ "supports_tool_choice": true }, "azure_ai/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -3318,6 +3329,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3364,6 +3376,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-2026-03-05": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3410,6 +3423,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3455,6 +3469,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro-2026-03-05": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3500,6 +3515,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3540,6 +3556,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-mini-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3580,6 +3597,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3620,6 +3638,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3849,6 +3868,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3918,6 +3938,7 @@ "supports_none_reasoning_effort": true }, "azure/eu/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3948,6 +3969,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -4107,6 +4129,7 @@ "supports_vision": true }, "azure/global-standard/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4155,6 +4178,7 @@ "supports_vision": true }, "azure/global/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4224,6 +4248,7 @@ "supports_none_reasoning_effort": true }, "azure/global/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4254,6 +4279,7 @@ "supports_vision": true }, "azure/global/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -4492,6 +4518,7 @@ "supports_vision": true }, "azure/gpt-4.1": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4559,6 +4586,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4626,6 +4654,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4902,6 +4931,7 @@ "supports_vision": false }, "azure/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", @@ -5344,6 +5374,7 @@ "supports_vision": true }, "azure/gpt-5": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5507,6 +5538,7 @@ "supports_vision": true }, "azure/gpt-5-mini": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5572,6 +5604,7 @@ "supports_vision": true }, "azure/gpt-5-nano": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 5e-09, "input_cost_per_token": 5e-08, "litellm_provider": "azure", @@ -5667,6 +5700,7 @@ "supports_vision": true }, "azure/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5736,6 +5770,7 @@ "supports_none_reasoning_effort": true }, "azure/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5797,6 +5832,7 @@ "supports_vision": true }, "azure/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5827,6 +5863,7 @@ "supports_vision": true }, "azure/gpt-5.2": { + "deprecation_date": "2027-06-08", "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", @@ -6136,6 +6173,7 @@ "supports_web_search": true }, "azure/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6180,6 +6218,7 @@ "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6218,6 +6257,7 @@ "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6379,6 +6419,7 @@ "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -7045,6 +7086,7 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -7095,6 +7137,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7142,6 +7185,7 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7408,6 +7452,7 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7489,6 +7534,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7601,6 +7647,7 @@ "output_cost_per_token": 0.0 }, "azure/high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7610,6 +7657,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7619,6 +7667,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7628,6 +7677,7 @@ ] }, "azure/low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7637,6 +7687,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7646,6 +7697,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7655,6 +7707,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7664,6 +7717,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7673,6 +7727,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7695,6 +7750,7 @@ ] }, "azure/gpt-image-1.5": { + "deprecation_date": "2027-06-16", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7720,6 +7776,7 @@ ] }, "azure/gpt-image-2": { + "deprecation_date": "2027-10-21", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7751,6 +7808,7 @@ "supports_pdf_input": true }, "azure/low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7760,6 +7818,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7769,6 +7828,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0345052083e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7778,6 +7838,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7787,6 +7848,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7796,6 +7858,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 7.9752604167e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7805,6 +7868,7 @@ ] }, "azure/high/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7814,6 +7878,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7823,6 +7888,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.1575520833e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7850,6 +7916,7 @@ "supports_function_calling": true }, "azure/o1": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -7944,6 +8011,7 @@ "supports_vision": false }, "azure/o3": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8041,6 +8109,7 @@ "supports_web_search": true }, "azure/o3-mini": { + "deprecation_date": "2026-10-01", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8071,6 +8140,7 @@ "supports_vision": false }, "azure/o3-pro": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8132,6 +8202,7 @@ "supports_vision": true }, "azure/o4-mini": { + "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8580,6 +8651,7 @@ "supports_vision": true }, "azure/us/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8649,6 +8721,7 @@ "supports_none_reasoning_effort": true }, "azure/us/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8679,6 +8752,7 @@ "supports_vision": true }, "azure/us/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -8876,6 +8950,7 @@ ] }, "azure_ai/FW-DeepSeek-V3.2": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, "input_cost_per_token": 6.2e-07, "litellm_provider": "azure_ai", @@ -8906,6 +8981,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure_ai", @@ -8921,6 +8997,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5.1": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.86e-07, "input_cost_per_token": 1.54e-06, "litellm_provider": "azure_ai", @@ -8987,6 +9064,7 @@ "supports_tool_choice": true }, "azure_ai/FW-Kimi-K2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 6.6e-07, "litellm_provider": "azure_ai", @@ -9079,6 +9157,7 @@ "supports_vision": true }, "azure_ai/FW-MiniMax-M2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, "input_cost_per_token": 3.3e-07, "litellm_provider": "azure_ai", @@ -9164,6 +9243,7 @@ ] }, "azure_ai/MAI-Image-2e": { + "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", "mode": "image_generation", @@ -9175,6 +9255,7 @@ ] }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9188,6 +9269,7 @@ "supports_vision": true }, "azure_ai/Llama-3.2-90B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 2.04e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9249,6 +9331,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-405B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 5.33e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9271,6 +9354,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9452,6 +9536,7 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2505": { + "deprecation_date": "2026-07-20", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", @@ -9529,6 +9614,7 @@ "output_cost_per_token": 0.0 }, "azure_ai/cohere-rerank-v3.5": { + "deprecation_date": "2026-05-14", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -9591,6 +9677,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-r1": { + "deprecation_date": "2026-08-13", "input_cost_per_token": 1.35e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9614,6 +9701,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3-0324": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.14e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9626,6 +9714,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.1": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.23e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9639,6 +9728,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.74e-06, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9652,6 +9742,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-flash": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.9e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9683,6 +9774,7 @@ "supports_embedding_image_input": true }, "azure_ai/global/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9697,6 +9789,7 @@ "supports_web_search": true }, "azure_ai/global/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9712,6 +9805,7 @@ "supports_web_search": true }, "azure_ai/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9726,6 +9820,7 @@ "supports_web_search": true }, "azure_ai/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9773,6 +9868,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-non-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9786,6 +9882,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9863,6 +9960,7 @@ "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { + "deprecation_date": "2027-01-26", "input_cost_per_token": 6e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -9877,6 +9975,7 @@ "supports_vision": true }, "azure_ai/kimi-k2.6": { + "deprecation_date": "2027-04-16", "input_cost_per_token": 9.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -10004,6 +10103,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -11999,6 +12099,7 @@ ] }, "claude-haiku-4-5-20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12022,6 +12123,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12170,6 +12272,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12203,6 +12306,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5-20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12237,6 +12341,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-5": { + "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -12273,6 +12378,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-6": { + "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12419,6 +12525,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12448,6 +12555,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12477,6 +12585,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12513,6 +12622,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6-20260205": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12549,6 +12659,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12587,6 +12698,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-opus-4-7-20260416": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12625,6 +12737,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { + "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -12660,6 +12773,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-5": { + "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12698,6 +12812,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-4-8": { + "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14801,6 +14916,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -18353,6 +18469,7 @@ } }, "gemini-2.5-flash": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18398,6 +18515,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18442,6 +18560,7 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -18522,6 +18641,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -18646,6 +18766,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -18702,6 +18823,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -18791,6 +18913,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, @@ -19062,6 +19185,7 @@ "supports_image_size": false }, "gemini-2.5-pro": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19373,6 +19497,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19809,6 +19934,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-robotics-er-1.6-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_audio_token": 2e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -19879,6 +20005,7 @@ "supports_vision": true }, "gemini-embedding-001": { + "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, @@ -21492,6 +21619,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -23004,6 +23132,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -24135,6 +24264,7 @@ "supports_pdf_input": true }, "low/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24146,6 +24276,7 @@ "supports_pdf_input": true }, "low/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24157,6 +24288,7 @@ "supports_pdf_input": true }, "low/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24168,6 +24300,7 @@ "supports_pdf_input": true }, "medium/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.034, "litellm_provider": "openai", "mode": "image_generation", @@ -24179,6 +24312,7 @@ "supports_pdf_input": true }, "medium/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24190,6 +24324,7 @@ "supports_pdf_input": true }, "medium/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24201,6 +24336,7 @@ "supports_pdf_input": true }, "high/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.133, "litellm_provider": "openai", "mode": "image_generation", @@ -24212,6 +24348,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24223,6 +24360,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24234,6 +24372,7 @@ "supports_pdf_input": true }, "standard/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24245,6 +24384,7 @@ "supports_pdf_input": true }, "standard/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24256,6 +24396,7 @@ "supports_pdf_input": true }, "standard/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24267,6 +24408,7 @@ "supports_pdf_input": true }, "1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24278,6 +24420,7 @@ "supports_pdf_input": true }, "1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24289,6 +24432,7 @@ "supports_pdf_input": true }, "1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -27202,18 +27346,21 @@ "output_cost_per_second": 0.0 }, "hd/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 7.629e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -27260,6 +27407,7 @@ "max_output_tokens": 8192 }, "high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.167, "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "openai", @@ -27270,6 +27418,7 @@ ] }, "high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -27280,6 +27429,7 @@ ] }, "high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -28067,6 +28217,7 @@ "supports_tool_choice": true }, "low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.011, "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "openai", @@ -28077,6 +28228,7 @@ ] }, "low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28087,6 +28239,7 @@ ] }, "low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28111,6 +28264,7 @@ "output_cost_per_image": 0.072 }, "medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.042, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28121,6 +28275,7 @@ ] }, "medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28131,6 +28286,7 @@ ] }, "medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28141,6 +28297,7 @@ ] }, "low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.005, "litellm_provider": "openai", "mode": "image_generation", @@ -28149,6 +28306,7 @@ ] }, "low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28157,6 +28315,7 @@ ] }, "low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28165,6 +28324,7 @@ ] }, "medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.011, "litellm_provider": "openai", "mode": "image_generation", @@ -28173,6 +28333,7 @@ ] }, "medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -28181,6 +28342,7 @@ ] }, "medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -30074,6 +30236,7 @@ ] }, "multimodalembedding@001": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2e-07, "input_cost_per_image": 0.0001, "input_cost_per_token": 8e-07, @@ -35772,18 +35935,21 @@ "output_cost_per_image": 0.14 }, "standard/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 3.81469e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -35847,6 +36013,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "text-embedding-005": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -35920,6 +36087,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35929,6 +36097,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35938,6 +36107,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35947,6 +36117,7 @@ "output_cost_per_token": 0.0 }, "text-multilingual-embedding-002": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -38434,6 +38605,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38457,6 +38629,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-haiku-4-5@20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38609,6 +38782,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38636,6 +38810,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38654,6 +38829,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-1@20250805": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38672,6 +38848,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38700,6 +38877,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-5@20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38729,6 +38907,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38759,6 +38938,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38789,6 +38969,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38820,6 +39001,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38851,6 +39033,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38882,6 +39065,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38913,6 +39097,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38945,6 +39130,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38977,6 +39163,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39009,6 +39196,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39041,6 +39229,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39069,6 +39258,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39131,6 +39321,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5@20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39160,6 +39351,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -39187,6 +39379,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39218,6 +39411,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39382,6 +39576,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -39427,6 +39622,7 @@ "supports_image_size": false }, "vertex_ai/gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -39459,6 +39655,7 @@ "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -39535,6 +39732,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -39591,6 +39789,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -40308,6 +40507,7 @@ "supports_tool_choice": true }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40322,6 +40522,7 @@ ] }, "vertex_ai/veo-3.0-fast-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40336,6 +40537,7 @@ ] }, "vertex_ai/veo-3.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40378,6 +40580,7 @@ ] }, "vertex_ai/veo-3.1-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40392,6 +40595,7 @@ ] }, "vertex_ai/veo-3.1-fast-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -46773,6 +46977,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, From 2adf8aa581745284a87ce08cf40fd659a471af3e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:08:00 -0700 Subject: [PATCH 18/88] feat(e2e): add record/replay transport seam and fixture bundle format E2E_FIXTURE_MODE selects the transport every e2e client is built on: live (default, unchanged behavior), record (pass through to the live proxy while writing every interaction to a fixture bundle), or replay (serve every interaction from the bundle with no proxy and no provider spend). Both new transports fulfil the existing Transport protocol, so no test changes shape. A bundle is a directory with a manifest (record timestamp, harness version, format version) and one JSON file per interaction, grouped per test in call order. Replay against a manifest older than seven days hard-fails at collection time naming the bundle age. Record always wipes and never reads the previous bundle, refusing to wipe a directory that is not a bundle. Auth header values are redacted on write; uploads store a sha256 digest. unique_marker() becomes deterministic per test in record/replay modes so a replay run regenerates exactly the requests the record run sent. Content-based match keys, streaming chunk fidelity, and provider-scoping are follow-ups (LIT-5741, LIT-5742, LIT-5745). --- .gitignore | 1 + tests/e2e/CLAUDE.md | 10 + tests/e2e/CONTRIBUTING.md | 11 + tests/e2e/conftest.py | 28 +- tests/e2e/e2e_config.py | 17 +- tests/e2e/fixture_bundle.py | 314 ++++++++++++++++ tests/e2e/fixture_transport.py | 550 ++++++++++++++++++++++++++++ tests/e2e/proxy_client.py | 37 +- tests/e2e/test_fixture_bundle.py | 218 +++++++++++ tests/e2e/test_fixture_transport.py | 438 ++++++++++++++++++++++ 10 files changed, 1609 insertions(+), 15 deletions(-) create mode 100644 tests/e2e/fixture_bundle.py create mode 100644 tests/e2e/fixture_transport.py create mode 100644 tests/e2e/test_fixture_bundle.py create mode 100644 tests/e2e/test_fixture_transport.py diff --git a/.gitignore b/.gitignore index 3329f39ca10..9b552a8c269 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .python-version .venv +tests/e2e/.fixtures/ .venv-typecheck .venv_policy_test .env diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 680e0dff67b..05753c736de 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -71,6 +71,16 @@ Request and response bodies are typed pydantic models in `models.py`; only the f Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself carries no marker and runs regardless. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +## Record and replay fixtures + +`E2E_FIXTURE_MODE` selects the transport every client is built on: `live` (the default, and what an unset variable means: nothing changes), `record` (run against the live proxy and write every interaction to a fixture bundle), or `replay` (serve every interaction back from the bundle with no HTTP at all, so a replay run needs no proxy and cannot bill a provider). The seam is `select_transport` in `fixture_transport.py`, applied inside `build_proxy_client`; both transports fulfil the same `Transport` protocol, so no test or client changes shape in any mode + +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format + +Replay matches calls per test by transport verb and path in recorded order and raises `ReplayMiss` on any drift, naming the recorded and the actual call; the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy + +Deliberately not here yet: canonical content-based match keys (LIT-5741), streaming chunk fidelity (LIT-5742), and scoping record/replay to provider-bound traffic (LIT-5745) + ## Typing The harness is fully typed with no error budget: `make lint-e2e-basedpyright` must report zero basedpyright errors, and CI enforces that on any PR touching `tests/e2e/**/*.py`. When a response field is untyped, model it in `models.py` (just the fields you read) and let pydantic validate it, rather than threading a `dict` or `Any` through the test diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index dc69bd42171..67da1be9562 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -52,6 +52,17 @@ The suites run against a live proxy, so bring one up first by running the litell Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy +### Record and replay + +`E2E_FIXTURE_MODE=record` runs a suite against the live proxy as usual while writing every request/response pair to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`); `E2E_FIXTURE_MODE=replay` then runs the same suite entirely from that bundle, with no proxy traffic and no provider spend; the proxy liveness gate is skipped, so replay runs with no proxy up at all. Unset (or `live`) behaves exactly as before the knob existed + +```bash +E2E_FIXTURE_MODE=record uv run pytest tests/e2e/llm_translation/ -v +E2E_FIXTURE_MODE=replay uv run pytest tests/e2e/llm_translation/ -v +``` + +Replay fails hard (`ReplayMiss`) when the tests drift from the recording, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. See `CLAUDE.md` in this directory for the bundle format and the transport seam + Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass ## What a complete test looks like diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index eff3b4ddf58..6b27bb459a5 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -16,12 +16,18 @@ shared fixtures build on it. import functools import os from collections.abc import Iterator +from datetime import datetime, timezone import pytest import requests -from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL +from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup +from fixture_transport import ( + fixture_mode_collection_error, + fixture_report_lines, + parse_fixture_mode, +) from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager from proxy_client import ProxyClient, build_proxy_client @@ -49,6 +55,21 @@ def pytest_configure(config: pytest.Config) -> None: ) +def pytest_sessionstart(session: pytest.Session) -> None: + """Abort before collection when E2E_FIXTURE_MODE can never work: an unknown + mode value, or replay against a missing, unreadable, or stale bundle (the + stale message names the bundle's age). Live and record modes pass through.""" + reason = fixture_mode_collection_error( + FIXTURE_MODE_RAW, FIXTURE_DIR, now=datetime.now(timezone.utc) + ) + if reason is not None: + raise pytest.UsageError(reason) + + +def pytest_report_header(config: pytest.Config) -> list[str]: + return fixture_report_lines(FIXTURE_MODE_RAW, FIXTURE_DIR, now=datetime.now(timezone.utc)) + + def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: """Attach the two custom signals (suite package and covered cell ids) to every test's user_properties so the standard JUnit report (`--junitxml`) records them @@ -91,9 +112,12 @@ def _proxy_fail_reason() -> str | None: def pytest_runtest_setup(item: pytest.Item) -> None: """Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe. Unmarked tests (unit coverage of the harness) don't touch the proxy, so they - run even when none is up. Never skip for a missing proxy.""" + run even when none is up. Never skip for a missing proxy. Replay mode serves + every call from the fixture bundle, so it needs no live proxy either.""" if item.get_closest_marker("e2e") is None: return + if parse_fixture_mode(FIXTURE_MODE_RAW) == "replay": + return reason = _proxy_fail_reason() if reason is not None: pytest.fail(reason) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 277478eebaf..a5c3729f4be 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -13,6 +13,8 @@ from pathlib import Path from dotenv import load_dotenv +from fixture_transport import deterministic_marker, parse_fixture_mode + # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). # Compose injects them into the proxy container, but pytest on the host does not # inherit that file unless we load it. override=False so a real shell export wins. @@ -90,6 +92,15 @@ PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15")) EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") +# Record/replay fixture selection (see fixture_transport.py). The raw mode value +# is parsed and validated there; "live" (the default, also for empty values) +# means the harness behaves exactly as before this knob existed. +FIXTURE_MODE_RAW = os.environ.get("E2E_FIXTURE_MODE", "live") +FIXTURE_DIR = Path( + os.environ.get("E2E_FIXTURE_DIR", "").strip() + or str(Path(__file__).resolve().parent / ".fixtures") +) + # Deliberately modest concurrency. The suite shares its proxy with every other # suite in the run, and 750 users at spawn rate 50 saturated the request path hard # enough to distort latency-sensitive neighbours (and to spend real provider money @@ -148,7 +159,11 @@ def datadog_mcp_url(*, toolsets: str = "core") -> str: def unique_marker() -> str: """A short unique token per call/run, so concurrent runs and the shared - response cache never collide on prompts, tags, or customer ids.""" + response cache never collide on prompts, tags, or customer ids. In record + and replay modes the token is deterministic per test instead, so a replay + run regenerates the exact requests the record run sent.""" + if parse_fixture_mode(FIXTURE_MODE_RAW) in ("record", "replay"): + return deterministic_marker() return uuid.uuid4().hex[:12] diff --git a/tests/e2e/fixture_bundle.py b/tests/e2e/fixture_bundle.py new file mode 100644 index 00000000000..5eff2cf2876 --- /dev/null +++ b/tests/e2e/fixture_bundle.py @@ -0,0 +1,314 @@ +"""On-disk fixture bundle format for record/replay e2e runs (LIT-5729). + +A bundle is a directory: one ``manifest.json`` (record timestamp + harness +version + format version) plus one subdirectory per test, holding one JSON file +per transport interaction in call order. Bundles older than +``MAX_BUNDLE_AGE`` hard-fail replay at collection time (see conftest), so a +green replay run can never certify against fixtures that have drifted more than +a week from the live proxy. + +This module owns the format only. The transports that produce and consume it +live in fixture_transport.py; canonical request matching, streaming chunk +fidelity, and provider-scoping are follow-ups (LIT-5741/5742/5745) and are +deliberately absent here, which is why every interaction file stores the full +redacted request even though replay today matches by call order. +""" + +from __future__ import annotations + +import hashlib +import re +import shutil +import subprocess +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Annotated, Final, Literal + +from pydantic import BaseModel, Field, JsonValue, TypeAdapter + +from e2e_http import ( + BinaryStream, + NetworkError, + ProbeResult, + RateLimitedError, + Result, + StreamingResponse, + Success, + UnauthorizedError, + UnknownApiError, + ValidationError, +) + +BUNDLE_FORMAT_VERSION: Final = 1 +MAX_BUNDLE_AGE: Final = timedelta(days=7) +MANIFEST_FILENAME: Final = "manifest.json" + +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +class Manifest(BaseModel): + format_version: int + recorded_at: datetime + harness_version: str + + +class RecordedRequest(BaseModel): + """The request as the transport saw it, auth header values redacted. + + Replay today only matches ``method`` (the transport verb, not the HTTP verb) + and ``path`` in call order; the rest is stored so LIT-5741 can move to + content-based match keys without re-recording. File uploads store a content + digest instead of the bytes.""" + + method: str + path: str + headers: dict[str, str] + params: dict[str, str] = {} + body: JsonValue | None = None + form: dict[str, str] | None = None + file_name: str | None = None + file_sha256: str | None = None + file_bytes: int | None = None + + +class RecordedResult(BaseModel): + """A ``Result[R]`` flattened for disk. ``data`` holds the success payload as + raw JSON; replay re-validates it against the ``response_type`` the caller + passes, exactly like a live response body.""" + + shape: Literal["result"] = "result" + kind: Literal["success", "network", "unauthorized", "rate_limited", "validation", "unknown"] + status_code: int | None = None + data: JsonValue | None = None + message: str | None = None + body: str | None = None + retry_after_seconds: int | None = None + + +class RecordedStreaming(BaseModel): + shape: Literal["streaming"] = "streaming" + payload: StreamingResponse + + +class RecordedBinary(BaseModel): + shape: Literal["binary"] = "binary" + payload: BinaryStream + + +class RecordedProbe(BaseModel): + shape: Literal["probe"] = "probe" + payload: ProbeResult + + +type RecordedResponse = RecordedResult | RecordedStreaming | RecordedBinary | RecordedProbe + + +class Interaction(BaseModel): + request: RecordedRequest + response: Annotated[ + RecordedResult | RecordedStreaming | RecordedBinary | RecordedProbe, + Field(discriminator="shape"), + ] + + +def to_json_value(model: BaseModel) -> JsonValue: + return _JSON.validate_json(model.model_dump_json(by_alias=True)) + + +def from_result[R: BaseModel](result: Result[R]) -> RecordedResult: + match result: + case Success(status_code=status_code, data=data): + return RecordedResult(kind="success", status_code=status_code, data=to_json_value(data)) + case NetworkError(message=message): + return RecordedResult(kind="network", message=message) + case UnauthorizedError(): + return RecordedResult(kind="unauthorized") + case RateLimitedError(retry_after_seconds=retry_after_seconds, body=body): + return RecordedResult(kind="rate_limited", retry_after_seconds=retry_after_seconds, body=body) + case ValidationError(message=message): + return RecordedResult(kind="validation", message=message) + case UnknownApiError(status_code=status_code, body=body): + return RecordedResult(kind="unknown", status_code=status_code, body=body) + + +def to_result[R: BaseModel](recorded: RecordedResult, response_type: type[R]) -> Result[R]: + match recorded.kind: + case "success": + return Success( + status_code=recorded.status_code or 200, + data=response_type.model_validate(recorded.data), + ) + case "network": + return NetworkError(message=recorded.message or "") + case "unauthorized": + return UnauthorizedError() + case "rate_limited": + return RateLimitedError( + retry_after_seconds=recorded.retry_after_seconds, body=recorded.body or "" + ) + case "validation": + return ValidationError(message=recorded.message or "") + case "unknown": + return UnknownApiError(status_code=recorded.status_code or 0, body=recorded.body or "") + + +def slugify(raw: str, *, limit: int = 60) -> str: + clean = re.sub(r"[^A-Za-z0-9_.-]+", "-", raw).strip("-") + return clean[:limit].rstrip("-") + + +def slug_for_test(test_key: str) -> str: + """Directory name for one test's interactions: a readable tail plus a short + digest of the full node id, so same-named methods in different classes or + files never collide.""" + digest = hashlib.sha1(test_key.encode()).hexdigest()[:8] + tail = slugify(test_key.rsplit("::", 1)[-1]) + return f"{tail}-{digest}" if tail else digest + + +def interaction_filename(ordinal: int, request: RecordedRequest) -> str: + path_part = slugify(request.path, limit=40) or "root" + return f"{ordinal:04d}-{request.method}-{path_part}.json" + + +def harness_version() -> str: + try: + proc = subprocess.run( + ("git", "rev-parse", "--short", "HEAD"), + cwd=Path(__file__).resolve().parent, + capture_output=True, + text=True, + timeout=10, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return "unknown" + return proc.stdout.strip() or "unknown" + + +@dataclass(slots=True) +class BundleRecorder: + """Appends interaction files under ``root``, one subdirectory per test, with + a per-test ordinal that fixes replay order. ``prepare_bundle`` is the only + constructor: it guarantees the directory started empty with a fresh + manifest, so record mode never reads (or merges into) an existing bundle.""" + + root: Path + _ordinals: dict[str, int] = field(default_factory=dict) + + def record(self, *, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None: + slug = slug_for_test(test_key) + ordinal = self._ordinals.get(slug, 0) + self._ordinals[slug] = ordinal + 1 + directory = self.root / slug + directory.mkdir(parents=True, exist_ok=True) + interaction = Interaction(request=request, response=response) + target = directory / interaction_filename(ordinal, request) + target.write_text(interaction.model_dump_json(indent=2), encoding="utf-8") + + +@dataclass(frozen=True, slots=True) +class UnsafeBundleDir: + path: Path + reason: str + + +def prepare_bundle(root: Path) -> BundleRecorder | UnsafeBundleDir: + """Start a fresh bundle at ``root`` for record mode: wipe whatever bundle is + there and write a new manifest. Refuses to wipe a directory that is neither + empty nor a bundle (no manifest.json), so a mistyped E2E_FIXTURE_DIR can + never delete unrelated files.""" + if root.exists(): + if not root.is_dir(): + return UnsafeBundleDir(path=root, reason="exists and is not a directory") + entries = tuple(root.iterdir()) + if entries and not (root / MANIFEST_FILENAME).is_file(): + return UnsafeBundleDir( + path=root, + reason=f"is not empty and has no {MANIFEST_FILENAME}; refusing to wipe a non-bundle directory", + ) + shutil.rmtree(root) + root.mkdir(parents=True) + manifest = Manifest( + format_version=BUNDLE_FORMAT_VERSION, + recorded_at=datetime.now(timezone.utc), + harness_version=harness_version(), + ) + (root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(indent=2), encoding="utf-8") + return BundleRecorder(root=root) + + +@dataclass(frozen=True, slots=True) +class FreshBundle: + manifest: Manifest + + +@dataclass(frozen=True, slots=True) +class StaleBundle: + recorded_at: datetime + age: timedelta + limit: timedelta + + +@dataclass(frozen=True, slots=True) +class UnreadableBundle: + reason: str + + +type BundleFreshness = FreshBundle | StaleBundle | UnreadableBundle + + +def _read_manifest(root: Path) -> Manifest | UnreadableBundle: + manifest_path = root / MANIFEST_FILENAME + if not manifest_path.is_file(): + return UnreadableBundle(reason=f"no {MANIFEST_FILENAME} found (record one with E2E_FIXTURE_MODE=record)") + try: + return Manifest.model_validate_json(manifest_path.read_text(encoding="utf-8")) + except ValueError as exc: + return UnreadableBundle(reason=f"{MANIFEST_FILENAME} is invalid: {exc}") + + +def check_freshness(root: Path, *, now: datetime) -> BundleFreshness: + manifest = _read_manifest(root) + if isinstance(manifest, UnreadableBundle): + return manifest + if manifest.format_version != BUNDLE_FORMAT_VERSION: + return UnreadableBundle( + reason=f"format_version {manifest.format_version} != supported {BUNDLE_FORMAT_VERSION}" + ) + recorded_at = ( + manifest.recorded_at + if manifest.recorded_at.tzinfo is not None + else manifest.recorded_at.replace(tzinfo=timezone.utc) + ) + age = now - recorded_at + if age > MAX_BUNDLE_AGE: + return StaleBundle(recorded_at=recorded_at, age=age, limit=MAX_BUNDLE_AGE) + return FreshBundle(manifest=manifest) + + +def format_age(age: timedelta) -> str: + total_hours = int(age.total_seconds()) // 3600 + return f"{total_hours // 24}d{total_hours % 24}h" + + +@dataclass(frozen=True, slots=True) +class LoadedBundle: + manifest: Manifest + interactions: dict[str, tuple[Interaction, ...]] + + +def load_bundle(root: Path) -> LoadedBundle | UnreadableBundle: + manifest = _read_manifest(root) + if isinstance(manifest, UnreadableBundle): + return manifest + interactions = { + directory.name: tuple( + Interaction.model_validate_json(file.read_text(encoding="utf-8")) + for file in sorted(directory.glob("*.json")) + ) + for directory in sorted(root.iterdir()) + if directory.is_dir() + } + return LoadedBundle(manifest=manifest, interactions=interactions) diff --git a/tests/e2e/fixture_transport.py b/tests/e2e/fixture_transport.py new file mode 100644 index 00000000000..756362cf29e --- /dev/null +++ b/tests/e2e/fixture_transport.py @@ -0,0 +1,550 @@ +"""Record/replay transports behind the same ``Transport`` protocol (LIT-5729). + +``RecordingTransport`` decorates the live transport: every call passes through +unchanged and its request/response pair is appended to the fixture bundle. +``ReplayTransport`` implements the protocol from a recorded bundle alone: no +HTTP, no proxy, no provider spend. Because both fulfil ``Transport``, no test +or client changes shape; ``build_proxy_client`` picks the transport from +``E2E_FIXTURE_MODE`` (live | record | replay, default live). + +Replay matches each call by test node id and call order, verifying transport +verb + path and failing hard on any drift (``ReplayMiss``). Canonical +content-based match keys are LIT-5741; streaming chunk fidelity is LIT-5742; +scoping record/replay to provider-bound traffic is LIT-5745. +""" + +from __future__ import annotations + +import functools +import hashlib +import os +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Final, Literal, assert_never + +from pydantic import BaseModel + +from e2e_http import AuthHeaders, BinaryStream, ProbeResult, Result, StreamingResponse +from fixture_bundle import ( + BundleRecorder, + FreshBundle, + Interaction, + LoadedBundle, + RecordedBinary, + RecordedProbe, + RecordedRequest, + RecordedResponse, + RecordedResult, + RecordedStreaming, + StaleBundle, + UnreadableBundle, + UnsafeBundleDir, + check_freshness, + format_age, + from_result, + load_bundle, + prepare_bundle, + slug_for_test, + to_json_value, + to_result, +) +from transport import Transport + +type FixtureMode = Literal["live", "record", "replay"] + +FIXTURE_MODES: Final[tuple[FixtureMode, ...]] = ("live", "record", "replay") + +SESSION_TEST_KEY: Final = "session" + +REDACTED_HEADER_NAMES: Final[frozenset[str]] = frozenset({"authorization", "x-litellm-api-key"}) +REDACTED_VALUE: Final = "" + + +@dataclass(frozen=True, slots=True) +class InvalidFixtureMode: + value: str + + +def parse_fixture_mode(raw: str) -> FixtureMode | InvalidFixtureMode: + normalized = raw.strip().lower() or "live" + match normalized: + case "live" | "record" | "replay": + return normalized + case _: + return InvalidFixtureMode(value=raw) + + +def current_test_key() -> str: + """The pytest node id of the running test, from the PYTEST_CURRENT_TEST env + var pytest maintains (`` (setup|call|teardown)``); ``session`` for + calls outside any test (e.g. session-finish cleanup).""" + raw = os.environ.get("PYTEST_CURRENT_TEST", "") + if not raw: + return SESSION_TEST_KEY + return raw.rsplit(" (", 1)[0] + + +class ReplayMiss(AssertionError): + """Replay had no recorded interaction for a call the suite made. The test + drifted from the bundle (or the bundle from the suite): re-record.""" + + +_marker_ordinals: Final[dict[str, int]] = {} + + +def deterministic_marker() -> str: + """Stable stand-in for uuid-based unique markers in record and replay modes: + the Nth marker of a test is a pure function of the test's node id and N, so a + replay run regenerates exactly the model names, prompts, and tags the record + run sent and every recorded poll response still satisfies its predicate.""" + test_key = current_test_key() + ordinal = _marker_ordinals.get(test_key, 0) + _marker_ordinals[test_key] = ordinal + 1 + return hashlib.sha1(f"{test_key}#{ordinal}".encode()).hexdigest()[:12] + + +def _dump_flat(model: BaseModel | None) -> dict[str, str]: + if model is None: + return {} + dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True) + return {key: str(value) for key, value in dumped.items()} + + +def _redact(headers: dict[str, str]) -> dict[str, str]: + return { + name: REDACTED_VALUE if name.lower() in REDACTED_HEADER_NAMES else value + for name, value in headers.items() + } + + +def recorded_request( + method: str, + path: str, + *, + headers: BaseModel, + body: BaseModel | None = None, + params: BaseModel | None = None, + form: BaseModel | None = None, + file_name: str | None = None, + file_content: bytes | None = None, +) -> RecordedRequest: + return RecordedRequest( + method=method, + path=path, + headers=_redact(_dump_flat(headers)), + params=_dump_flat(params), + body=None if body is None else to_json_value(body), + form=None if form is None else _dump_flat(form), + file_name=file_name, + file_sha256=None if file_content is None else hashlib.sha256(file_content).hexdigest(), + file_bytes=None if file_content is None else len(file_content), + ) + + +@dataclass(frozen=True, slots=True) +class RecordingTransport: + """Decorator over the live transport: forwards every call and appends the + interaction to the bundle, so a green live run leaves behind exactly the + traffic replay needs.""" + + inner: Transport + recorder: BundleRecorder + + def _record(self, request: RecordedRequest, response: RecordedResponse) -> None: + self.recorder.record(test_key=current_test_key(), request=request, response=response) + + def bearer(self, key: str) -> AuthHeaders: + return self.inner.bearer(key) + + @property + def master(self) -> AuthHeaders: + return self.inner.master + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + result = self.inner.post(path, headers=headers, json=json, response_type=response_type) + self._record(recorded_request("post", path, headers=headers, body=json), from_result(result)) + return result + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + timeout: float | None = None, + ) -> Result[R]: + result = self.inner.get( + path, headers=headers, params=params, response_type=response_type, timeout=timeout + ) + self._record(recorded_request("get", path, headers=headers, params=params), from_result(result)) + return result + + def delete[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, + ) -> Result[R]: + result = self.inner.delete( + path, headers=headers, json=json, response_type=response_type, params=params + ) + self._record( + recorded_request("delete", path, headers=headers, body=json, params=params), + from_result(result), + ) + return result + + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + result = self.inner.patch(path, headers=headers, json=json, response_type=response_type) + self._record(recorded_request("patch", path, headers=headers, body=json), from_result(result)) + return result + + def put[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + result = self.inner.put(path, headers=headers, json=json, response_type=response_type) + self._record(recorded_request("put", path, headers=headers, body=json), from_result(result)) + return result + + def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: + response = self.inner.stream(path, headers=headers, json=json) + self._record( + recorded_request("stream", path, headers=headers, body=json), + RecordedStreaming(payload=response), + ) + return response + + def stream_binary( + self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 + ) -> BinaryStream: + response = self.inner.stream_binary(path, headers=headers, json=json, chunk_size=chunk_size) + self._record( + recorded_request("stream_binary", path, headers=headers, body=json), + RecordedBinary(payload=response), + ) + return response + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + response = self.inner.send(path, headers=headers, json=json, params=params, stream=stream) + self._record( + recorded_request("send", path, headers=headers, body=json, params=params), + RecordedStreaming(payload=response), + ) + return response + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + response = self.inner.probe(path, params=params) + self._record( + recorded_request("probe", path, headers=self.master, params=params), + RecordedProbe(payload=response), + ) + return response + + def upload[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + form: BaseModel, + filename: str, + content: bytes, + file_content_type: str = "application/jsonl", + file_field: str = "file", + params: BaseModel | None = None, + response_type: type[R], + ) -> Result[R]: + result = self.inner.upload( + path, + headers=headers, + form=form, + filename=filename, + content=content, + file_content_type=file_content_type, + file_field=file_field, + params=params, + response_type=response_type, + ) + self._record( + recorded_request( + "upload", + path, + headers=headers, + params=params, + form=form, + file_name=filename, + file_content=content, + ), + from_result(result), + ) + return result + + def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: + response = self.inner.download(path, headers=headers) + self._record( + recorded_request("download", path, headers=headers), + RecordedStreaming(payload=response), + ) + return response + + +@dataclass(slots=True) +class ReplaySource: + """One shared cursor set over a loaded bundle, so every client built in the + session consumes the same recorded sequence per test.""" + + bundle: LoadedBundle + _cursors: dict[str, int] = field(default_factory=dict) + + def next_interaction(self, method: str, path: str) -> Interaction: + test_key = current_test_key() + slug = slug_for_test(test_key) + recorded = self.bundle.interactions.get(slug, ()) + index = self._cursors.get(slug, 0) + if index >= len(recorded): + raise ReplayMiss( + f"replay exhausted for {test_key}: call #{index + 1} ({method} {path}) has no recorded " + f"interaction ({len(recorded)} recorded under {slug}); re-record with E2E_FIXTURE_MODE=record" + ) + interaction = recorded[index] + if interaction.request.method != method or interaction.request.path != path: + raise ReplayMiss( + f"replay mismatch for {test_key} at call #{index + 1}: recorded " + f"{interaction.request.method} {interaction.request.path}, test made {method} {path}; " + "re-record with E2E_FIXTURE_MODE=record" + ) + self._cursors[slug] = index + 1 + return interaction + + +def _expect_result(interaction: Interaction) -> RecordedResult: + match interaction.response: + case RecordedResult() as recorded: + return recorded + case RecordedStreaming() | RecordedBinary() | RecordedProbe(): + raise ReplayMiss( + f"recorded {interaction.request.method} {interaction.request.path} is not a typed result" + ) + + +def _expect_streaming(interaction: Interaction) -> StreamingResponse: + match interaction.response: + case RecordedStreaming(payload=payload): + return payload + case RecordedResult() | RecordedBinary() | RecordedProbe(): + raise ReplayMiss( + f"recorded {interaction.request.method} {interaction.request.path} is not a streaming response" + ) + + +@dataclass(frozen=True, slots=True) +class ReplayTransport: + """A ``Transport`` served entirely from a recorded bundle: never opens a + connection, so a replay run cannot bill a provider.""" + + source: ReplaySource + master_key: str + + def bearer(self, key: str) -> AuthHeaders: + return AuthHeaders(authorization=f"Bearer {key}") + + @property + def master(self) -> AuthHeaders: + return self.bearer(self.master_key) + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("post", path)), response_type) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + timeout: float | None = None, + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("get", path)), response_type) + + def delete[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("delete", path)), response_type) + + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("patch", path)), response_type) + + def put[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("put", path)), response_type) + + def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: + return _expect_streaming(self.source.next_interaction("stream", path)) + + def stream_binary( + self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 + ) -> BinaryStream: + interaction = self.source.next_interaction("stream_binary", path) + match interaction.response: + case RecordedBinary(payload=payload): + return payload + case RecordedResult() | RecordedStreaming() | RecordedProbe(): + raise ReplayMiss( + f"recorded stream_binary {interaction.request.path} is not a binary stream" + ) + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + return _expect_streaming(self.source.next_interaction("send", path)) + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + interaction = self.source.next_interaction("probe", path) + match interaction.response: + case RecordedProbe(payload=payload): + return payload + case RecordedResult() | RecordedStreaming() | RecordedBinary(): + raise ReplayMiss(f"recorded probe {interaction.request.path} is not a probe result") + + def upload[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + form: BaseModel, + filename: str, + content: bytes, + file_content_type: str = "application/jsonl", + file_field: str = "file", + params: BaseModel | None = None, + response_type: type[R], + ) -> Result[R]: + return to_result(_expect_result(self.source.next_interaction("upload", path)), response_type) + + def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: + return _expect_streaming(self.source.next_interaction("download", path)) + + +@functools.lru_cache(maxsize=8) +def _shared_recorder(root: Path) -> BundleRecorder: + prepared = prepare_bundle(root) + if isinstance(prepared, UnsafeBundleDir): + raise ValueError(f"E2E_FIXTURE_DIR {prepared.path} {prepared.reason}") + return prepared + + +@functools.lru_cache(maxsize=8) +def _shared_replay_source(root: Path) -> ReplaySource: + loaded = load_bundle(root) + if isinstance(loaded, UnreadableBundle): + raise ValueError(f"cannot replay from {root}: {loaded.reason}") + return ReplaySource(bundle=loaded) + + +def select_transport( + live: Transport, *, mode_raw: str, bundle_dir: Path, master_key: str +) -> Transport: + """The one seam every client build goes through: wraps (record), replaces + (replay), or passes through (live) the transport per E2E_FIXTURE_MODE. The + recorder and replay cursors are process-wide singletons per bundle dir, so + every client in a session shares one bundle and one recorded sequence.""" + mode = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode(value=value): + raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}") + case "live": + return live + case "record": + return RecordingTransport(inner=live, recorder=_shared_recorder(bundle_dir)) + case "replay": + return ReplayTransport(source=_shared_replay_source(bundle_dir), master_key=master_key) + case _: + assert_never(mode) + + +def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datetime) -> str | None: + """Session-abort reason for a fixture-mode setup that can never work, or None. + Called at collection time (conftest pytest_sessionstart) so a stale or missing + bundle fails the whole run up front, naming the bundle age, instead of failing + every test individually.""" + mode = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode(value=value): + return f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}" + case "live" | "record": + return None + case "replay": + freshness = check_freshness(bundle_dir, now=now) + match freshness: + case FreshBundle(): + return None + case StaleBundle(recorded_at=recorded_at, age=age, limit=limit): + return ( + f"fixture bundle at {bundle_dir} is stale: recorded {recorded_at.isoformat()}, " + f"age {format_age(age)} exceeds the {limit.days}-day limit; " + "re-record with E2E_FIXTURE_MODE=record" + ) + case UnreadableBundle(reason=reason): + return f"E2E_FIXTURE_MODE=replay cannot use bundle at {bundle_dir}: {reason}" + case _: + assert_never(freshness) + case _: + assert_never(mode) + + +def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> list[str]: + """pytest report-header lines; empty in live mode so an unset + E2E_FIXTURE_MODE keeps today's output byte-identical.""" + mode = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode() | "live": + return [] + case "record": + return [f"e2e fixture mode: record -> {bundle_dir}"] + case "replay": + freshness = check_freshness(bundle_dir, now=now) + match freshness: + case FreshBundle(manifest=manifest): + return [ + f"e2e fixture mode: replay <- {bundle_dir} " + f"(recorded {manifest.recorded_at.isoformat()}, harness {manifest.harness_version})" + ] + case StaleBundle() | UnreadableBundle(): + return [f"e2e fixture mode: replay <- {bundle_dir}"] + case _: + assert_never(freshness) + case _: + assert_never(mode) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 5050b6fce68..843799ede6c 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -65,6 +65,8 @@ from models import ( ) from e2e_config import ( CONTROL_PLANE_BASE_URL, + FIXTURE_DIR, + FIXTURE_MODE_RAW, MASTER_KEY, POLL_INTERVAL, POLL_TIMEOUT, @@ -72,6 +74,7 @@ from e2e_config import ( REQUEST_TIMEOUT, settle_propagation, ) +from fixture_transport import select_transport from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -531,19 +534,29 @@ def build_proxy_client( The endpoints are injectable for callers that resolve the proxy some other way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must pass all three together, since a caller that overrides only the data plane - would leave management calls pointed at the env default.""" + would leave management calls pointed at the env default. + + E2E_FIXTURE_MODE wraps (record) or replaces (replay) the transport here, so + every client built from this seam records or replays without changing shape; + unset it stays the plain SplitTransport (see fixture_transport.py).""" + split = SplitTransport( + data=HttpTransport( + base_url=base_url, + master_key=master_key, + request_timeout=REQUEST_TIMEOUT, + ), + control=HttpTransport( + base_url=control_plane_base_url, + master_key=master_key, + request_timeout=REQUEST_TIMEOUT, + ), + ) return ProxyClient( - transport=SplitTransport( - data=HttpTransport( - base_url=base_url, - master_key=master_key, - request_timeout=REQUEST_TIMEOUT, - ), - control=HttpTransport( - base_url=control_plane_base_url, - master_key=master_key, - request_timeout=REQUEST_TIMEOUT, - ), + transport=select_transport( + split, + mode_raw=FIXTURE_MODE_RAW, + bundle_dir=FIXTURE_DIR, + master_key=master_key, ), poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, diff --git a/tests/e2e/test_fixture_bundle.py b/tests/e2e/test_fixture_bundle.py new file mode 100644 index 00000000000..fd4cca6451f --- /dev/null +++ b/tests/e2e/test_fixture_bundle.py @@ -0,0 +1,218 @@ +"""Harness coverage for the on-disk fixture bundle format (LIT-5729). + +No proxy and no ``e2e`` marker: these pin the bundle CONTRACT - the seven-day +freshness gate that names the bundle's age, record mode's wipe safety (never +delete a directory that is not a bundle), collision-free per-test slugs, and +lossless Result round-trips - so replay can never silently drift from what +record wrote. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from e2e_http import ( + NetworkError, + RateLimitedError, + Result, + Success, + UnauthorizedError, + UnknownApiError, + ValidationError, +) +from fixture_bundle import ( + BUNDLE_FORMAT_VERSION, + MANIFEST_FILENAME, + MAX_BUNDLE_AGE, + BundleRecorder, + FreshBundle, + LoadedBundle, + Manifest, + RecordedRequest, + RecordedResult, + StaleBundle, + UnreadableBundle, + UnsafeBundleDir, + check_freshness, + format_age, + from_result, + interaction_filename, + load_bundle, + prepare_bundle, + slug_for_test, + to_result, +) + +NOW = datetime(2026, 8, 18, 12, 0, 0, tzinfo=timezone.utc) + + +class Payload(BaseModel): + value: str + + +def write_manifest( + root: Path, recorded_at: datetime, *, format_version: int = BUNDLE_FORMAT_VERSION +) -> None: + root.mkdir(parents=True, exist_ok=True) + manifest = Manifest( + format_version=format_version, recorded_at=recorded_at, harness_version="abc1234" + ) + (root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(), encoding="utf-8") + + +def prepared(root: Path) -> BundleRecorder: + recorder = prepare_bundle(root) + assert isinstance(recorder, BundleRecorder) + return recorder + + +def plain_request(path: str) -> RecordedRequest: + return RecordedRequest(method="post", path=path, headers={}) + + +class TestResultRoundTrip: + @pytest.mark.parametrize( + "result", + [ + Success(status_code=201, data=Payload(value="ok")), + NetworkError(message="connection refused"), + UnauthorizedError(), + RateLimitedError(retry_after_seconds=7, body="slow down"), + ValidationError(message="bad shape"), + UnknownApiError(status_code=502, body="upstream exploded"), + ], + ) + def test_every_result_kind_survives_disk_and_back(self, result: Result[Payload]) -> None: + assert to_result(from_result(result), Payload) == result + + +class TestFreshness: + def test_bundle_at_the_limit_is_still_fresh(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - MAX_BUNDLE_AGE) + assert isinstance(check_freshness(root, now=NOW), FreshBundle) + + def test_stale_bundle_reports_age_and_limit(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=8, hours=3)) + freshness = check_freshness(root, now=NOW) + assert isinstance(freshness, StaleBundle) + assert freshness.age == timedelta(days=8, hours=3) + assert format_age(freshness.age) == "8d3h" + assert freshness.limit == MAX_BUNDLE_AGE + + def test_naive_recorded_at_is_read_as_utc(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, (NOW - timedelta(days=1)).replace(tzinfo=None)) + assert isinstance(check_freshness(root, now=NOW), FreshBundle) + + def test_missing_manifest_is_unreadable_with_recording_hint(self, tmp_path: Path) -> None: + freshness = check_freshness(tmp_path / "absent", now=NOW) + assert isinstance(freshness, UnreadableBundle) + assert MANIFEST_FILENAME in freshness.reason + assert "E2E_FIXTURE_MODE=record" in freshness.reason + + def test_corrupt_manifest_is_unreadable(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + root.mkdir() + (root / MANIFEST_FILENAME).write_text("{not json", encoding="utf-8") + assert isinstance(check_freshness(root, now=NOW), UnreadableBundle) + + def test_unknown_format_version_is_unreadable(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW, format_version=BUNDLE_FORMAT_VERSION + 1) + freshness = check_freshness(root, now=NOW) + assert isinstance(freshness, UnreadableBundle) + assert f"format_version {BUNDLE_FORMAT_VERSION + 1}" in freshness.reason + + +class TestPrepareBundle: + def test_fresh_directory_gets_a_fresh_manifest(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + prepared(root) + freshness = check_freshness(root, now=datetime.now(timezone.utc)) + assert isinstance(freshness, FreshBundle) + assert freshness.manifest.format_version == BUNDLE_FORMAT_VERSION + assert freshness.manifest.harness_version + + def test_record_wipes_the_previous_bundle_instead_of_reading_it(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + prepared(root).record( + test_key="old.py::test_old", + request=plain_request("/stale"), + response=RecordedResult(kind="unauthorized"), + ) + assert any(entry.is_dir() for entry in root.iterdir()) + prepared(root) + assert {entry.name for entry in root.iterdir()} == {MANIFEST_FILENAME} + + def test_refuses_to_wipe_a_directory_that_is_not_a_bundle(self, tmp_path: Path) -> None: + root = tmp_path / "precious" + root.mkdir() + (root / "notes.txt").write_text("keep me", encoding="utf-8") + outcome = prepare_bundle(root) + assert isinstance(outcome, UnsafeBundleDir) + assert MANIFEST_FILENAME in outcome.reason + assert (root / "notes.txt").read_text(encoding="utf-8") == "keep me" + + def test_refuses_a_path_that_is_a_file(self, tmp_path: Path) -> None: + target = tmp_path / "not-a-dir" + target.write_text("x", encoding="utf-8") + outcome = prepare_bundle(target) + assert isinstance(outcome, UnsafeBundleDir) + assert "not a directory" in outcome.reason + + +class TestSlugs: + def test_slug_for_test_is_deterministic(self) -> None: + key = "tests/e2e/suite/test_mod.py::TestX::test_case" + assert slug_for_test(key) == slug_for_test(key) + + def test_same_tail_in_different_files_never_collides(self) -> None: + first = slug_for_test("tests/e2e/a/test_a.py::test_case") + second = slug_for_test("tests/e2e/b/test_b.py::test_case") + assert first != second + assert first.startswith("test_case-") + assert second.startswith("test_case-") + + def test_interaction_filename_orders_and_slugs(self) -> None: + request = RecordedRequest(method="post", path="/chat/completions", headers={}) + assert interaction_filename(3, request) == "0003-post-chat-completions.json" + + +class TestRecordAndLoad: + def test_load_returns_interactions_in_recorded_order(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorder = prepared(root) + key = "suite/test_mod.py::test_ordered" + for path in ("/first", "/second", "/third"): + recorder.record( + test_key=key, + request=plain_request(path), + response=RecordedResult(kind="unauthorized"), + ) + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + assert [ + interaction.request.path for interaction in loaded.interactions[slug_for_test(key)] + ] == ["/first", "/second", "/third"] + + def test_interactions_group_per_test(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorder = prepared(root) + for key in ("suite/test_a.py::test_one", "suite/test_b.py::test_two"): + recorder.record( + test_key=key, + request=plain_request(f"/{key[-3:]}"), + response=RecordedResult(kind="unauthorized"), + ) + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + assert set(loaded.interactions) == { + slug_for_test("suite/test_a.py::test_one"), + slug_for_test("suite/test_b.py::test_two"), + } diff --git a/tests/e2e/test_fixture_transport.py b/tests/e2e/test_fixture_transport.py new file mode 100644 index 00000000000..6ffcdaeb95f --- /dev/null +++ b/tests/e2e/test_fixture_transport.py @@ -0,0 +1,438 @@ +"""Harness coverage for the record/replay transports (LIT-5729). + +No proxy and no ``e2e`` marker. A fake in-memory ``Transport`` stands in for +the live one (dependency injection, no monkeypatching): recording must pass +every value through unchanged while writing one redacted interaction file per +call, and replay must serve identical values from the bundle alone - the +fake's call log proves nothing reaches the inner transport - failing hard +(``ReplayMiss``) on any drift in order, verb, or path. The collection-time +gate and report header are pinned here too, including the stale message that +names the bundle's age. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from e2e_http import ( + AuthHeaders, + BinaryStream, + ProbeResult, + Result, + StreamingResponse, + Success, +) +from fixture_bundle import ( + BUNDLE_FORMAT_VERSION, + MANIFEST_FILENAME, + BundleRecorder, + Interaction, + LoadedBundle, + Manifest, + load_bundle, + prepare_bundle, + slug_for_test, +) +from fixture_transport import ( + InvalidFixtureMode, + RecordingTransport, + ReplayMiss, + ReplaySource, + ReplayTransport, + current_test_key, + deterministic_marker, + fixture_mode_collection_error, + fixture_report_lines, + parse_fixture_mode, + select_transport, +) +from transport import Transport + +NOW = datetime(2026, 8, 18, 12, 0, 0, tzinfo=timezone.utc) + + +class Payload(BaseModel): + value: str + + +class Body(BaseModel): + prompt: str + + +class Query(BaseModel): + q: str + + +STREAMING = StreamingResponse( + status_code=200, + body="", + content_type="text/event-stream", + chunks=2, + stream_events=["one", "two"], + stream_done=True, +) +BINARY = BinaryStream(status_code=200, content_type="audio/mpeg", chunk_count=3, total_bytes=42) +PROBE = ProbeResult(status_code=200, body="alive") + + +@dataclass +class FakeTransport: + calls: list[str] = field(default_factory=list) + + def bearer(self, key: str) -> AuthHeaders: + return AuthHeaders(authorization=f"Bearer {key}") + + @property + def master(self) -> AuthHeaders: + return self.bearer("sk-fake-master") + + def _success[R: BaseModel](self, response_type: type[R]) -> Result[R]: + return Success(status_code=200, data=response_type.model_validate({"value": "live"})) + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + self.calls.append(f"post {path}") + return self._success(response_type) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + timeout: float | None = None, + ) -> Result[R]: + self.calls.append(f"get {path}") + return self._success(response_type) + + def delete[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, + ) -> Result[R]: + self.calls.append(f"delete {path}") + return self._success(response_type) + + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + self.calls.append(f"patch {path}") + return self._success(response_type) + + def put[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + self.calls.append(f"put {path}") + return self._success(response_type) + + def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: + self.calls.append(f"stream {path}") + return STREAMING + + def stream_binary( + self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 + ) -> BinaryStream: + self.calls.append(f"stream_binary {path}") + return BINARY + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + self.calls.append(f"send {path}") + return STREAMING + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + self.calls.append(f"probe {path}") + return PROBE + + def upload[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + form: BaseModel, + filename: str, + content: bytes, + file_content_type: str = "application/jsonl", + file_field: str = "file", + params: BaseModel | None = None, + response_type: type[R], + ) -> Result[R]: + self.calls.append(f"upload {path}") + return self._success(response_type) + + def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: + self.calls.append(f"download {path}") + return STREAMING + + +def make_recorder(root: Path) -> BundleRecorder: + recorder = prepare_bundle(root) + assert isinstance(recorder, BundleRecorder) + return recorder + + +def replay_source(root: Path) -> ReplaySource: + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + return ReplaySource(bundle=loaded) + + +def this_tests_files(root: Path) -> list[Path]: + slug_dir = root / slug_for_test(current_test_key()) + return sorted(slug_dir.glob("*.json")) if slug_dir.is_dir() else [] + + +def write_manifest(root: Path, recorded_at: datetime) -> None: + root.mkdir(parents=True, exist_ok=True) + manifest = Manifest( + format_version=BUNDLE_FORMAT_VERSION, recorded_at=recorded_at, harness_version="abc1234" + ) + (root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(), encoding="utf-8") + + +class TestParseFixtureMode: + @pytest.mark.parametrize( + ("raw", "expected"), + [("live", "live"), ("record", "record"), ("replay", "replay"), ("", "live"), (" REPLAY ", "replay")], + ) + def test_known_values_normalize(self, raw: str, expected: str) -> None: + assert parse_fixture_mode(raw) == expected + + def test_unknown_value_is_invalid_with_the_original_spelling(self) -> None: + assert parse_fixture_mode("cached") == InvalidFixtureMode(value="cached") + + +class TestDeterministicMarker: + def test_sequence_is_a_pure_function_of_test_and_ordinal(self) -> None: + """A replay process must regenerate exactly the markers the record + process generated, so the Nth marker of a test is pinned to a pure + function of the node id and N.""" + key = current_test_key() + assert deterministic_marker() == hashlib.sha1(f"{key}#0".encode()).hexdigest()[:12] + assert deterministic_marker() == hashlib.sha1(f"{key}#1".encode()).hexdigest()[:12] + + +class TestCurrentTestKey: + def test_names_this_test_and_strips_the_phase(self) -> None: + key = current_test_key() + assert key.endswith("TestCurrentTestKey::test_names_this_test_and_strips_the_phase") + assert "(call)" not in key + + +class TestRecordingTransport: + def test_passes_the_result_through_and_writes_one_file_per_call(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + result = recording.post( + "/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload + ) + assert result == Success(status_code=200, data=Payload(value="live")) + assert fake.calls == ["post /model/new"] + files = this_tests_files(root) + assert [file.name for file in files] == ["0000-post-model-new.json"] + interaction = Interaction.model_validate_json(files[0].read_text(encoding="utf-8")) + assert interaction.request.method == "post" + assert interaction.request.path == "/model/new" + + def test_redacts_auth_header_values_in_the_recorded_request(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + headers = AuthHeaders.model_validate( + {"authorization": "Bearer sk-secret", "x-litellm-api-key": "sk-other"} + ) + recording.post("/key/generate", headers=headers, json=Body(prompt="x"), response_type=Payload) + interaction = Interaction.model_validate_json( + this_tests_files(root)[0].read_text(encoding="utf-8") + ) + assert interaction.request.headers == { + "authorization": "", + "x-litellm-api-key": "", + } + assert "sk-secret" not in this_tests_files(root)[0].read_text(encoding="utf-8") + + def test_upload_records_a_content_digest_not_the_bytes(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.upload( + "/v1/files", + headers=fake.master, + form=Query(q="batch"), + filename="batch.jsonl", + content=b'{"custom_id": "1"}', + response_type=Payload, + ) + interaction = Interaction.model_validate_json( + this_tests_files(root)[0].read_text(encoding="utf-8") + ) + assert interaction.request.file_name == "batch.jsonl" + assert interaction.request.file_bytes == len(b'{"custom_id": "1"}') + assert interaction.request.file_sha256 is not None + assert "custom_id" not in interaction.request.model_dump_json() + + +class TestReplayTransport: + def test_serves_recorded_values_without_touching_the_inner_transport( + self, tmp_path: Path + ) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recorded_post = recording.post( + "/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload + ) + recorded_get = recording.get( + "/v1/models", headers=fake.master, params=Query(q="all"), response_type=Payload + ) + recorded_stream = recording.stream( + "/chat/completions", headers=fake.master, json=Body(prompt="hi") + ) + recorded_probe = recording.probe("/health/liveliness", params=Query(q="1")) + recorded_binary = recording.stream_binary( + "/v1/audio/speech", headers=fake.master, json=Body(prompt="say") + ) + calls_after_record = list(fake.calls) + + replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") + assert ( + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + == recorded_post + ) + assert ( + replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) + == recorded_get + ) + assert ( + replay.stream("/chat/completions", headers=replay.master, json=Body(prompt="hi")) + == recorded_stream + ) + assert replay.probe("/health/liveliness", params=Query(q="1")) == recorded_probe + assert ( + replay.stream_binary("/v1/audio/speech", headers=replay.master, json=Body(prompt="say")) + == recorded_binary + ) + assert fake.calls == calls_after_record + + def test_mismatched_call_names_recorded_and_actual(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") + with pytest.raises(ReplayMiss, match=r"recorded post /model/new, test made get /v1/models"): + replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) + + def test_exhausted_recording_names_the_call_count(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + with pytest.raises(ReplayMiss, match=r"call #2 \(post /model/new\) has no recorded interaction \(1 recorded"): + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + + +class TestSelectTransport: + def test_live_returns_the_live_transport_untouched(self, tmp_path: Path) -> None: + fake = FakeTransport() + for mode_raw in ("live", ""): + assert ( + select_transport(fake, mode_raw=mode_raw, bundle_dir=tmp_path / "b", master_key="sk") + is fake + ) + + def test_record_wraps_live_and_starts_a_fresh_bundle(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=30)) + (root / "old-test-slug").mkdir() + (root / "old-test-slug" / "0000-post-old.json").write_text("{}", encoding="utf-8") + selected = select_transport(fake, mode_raw="record", bundle_dir=root, master_key="sk") + assert isinstance(selected, RecordingTransport) + assert selected.inner is fake + assert {entry.name for entry in root.iterdir()} == {MANIFEST_FILENAME} + + def test_replay_builds_a_transport_from_the_bundle_alone(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + make_recorder(root) + selected = select_transport(fake, mode_raw="replay", bundle_dir=root, master_key="sk-master") + assert isinstance(selected, ReplayTransport) + assert selected.master == AuthHeaders(authorization="Bearer sk-master") + + def test_invalid_mode_raises_naming_the_value(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="cached"): + select_transport( + FakeTransport(), mode_raw="cached", bundle_dir=tmp_path / "b", master_key="sk" + ) + + +class TestCollectionGate: + def test_invalid_mode_names_the_value_and_the_choices(self, tmp_path: Path) -> None: + assert ( + fixture_mode_collection_error("cached", tmp_path, now=NOW) + == "E2E_FIXTURE_MODE='cached' is not one of live, record, replay" + ) + + @pytest.mark.parametrize("mode_raw", ["live", "", "record"]) + def test_live_and_record_never_block_collection(self, mode_raw: str, tmp_path: Path) -> None: + assert fixture_mode_collection_error(mode_raw, tmp_path / "missing", now=NOW) is None + + def test_replay_with_no_bundle_says_how_to_record_one(self, tmp_path: Path) -> None: + reason = fixture_mode_collection_error("replay", tmp_path / "missing", now=NOW) + assert reason is not None + assert f"no {MANIFEST_FILENAME}" in reason + assert "E2E_FIXTURE_MODE=record" in reason + + def test_stale_replay_bundle_fails_naming_its_age(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=9, hours=5)) + reason = fixture_mode_collection_error("replay", root, now=NOW) + assert reason is not None + assert "age 9d5h exceeds the 7-day limit" in reason + assert "re-record with E2E_FIXTURE_MODE=record" in reason + + def test_fresh_replay_bundle_collects(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=2)) + assert fixture_mode_collection_error("replay", root, now=NOW) is None + + +class TestReportHeader: + def test_live_mode_prints_nothing(self, tmp_path: Path) -> None: + assert fixture_report_lines("live", tmp_path, now=NOW) == [] + assert fixture_report_lines("", tmp_path, now=NOW) == [] + + def test_record_and_replay_name_the_bundle(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorded_at = NOW - timedelta(days=1) + write_manifest(root, recorded_at) + assert fixture_report_lines("record", root, now=NOW) == [ + f"e2e fixture mode: record -> {root}" + ] + replay_lines = fixture_report_lines("replay", root, now=NOW) + assert len(replay_lines) == 1 + assert "replay" in replay_lines[0] + assert recorded_at.isoformat() in replay_lines[0] From 6bf535bb8f99754bd9524da840ad7b350047363f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:34:23 -0700 Subject: [PATCH 19/88] feat(e2e): fail passed replays that leave recorded interactions unconsumed --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/conftest.py | 34 ++++++++++++++++++++- tests/e2e/fixture_transport.py | 24 +++++++++++++++ tests/e2e/test_fixture_transport.py | 47 +++++++++++++++++++++++++++++ 4 files changed, 105 insertions(+), 2 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 05753c736de..9969ed10308 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -77,7 +77,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format -Replay matches calls per test by transport verb and path in recorded order and raises `ReplayMiss` on any drift, naming the recorded and the actual call; the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy +Replay matches calls per test by transport verb and path in recorded order and raises `ReplayMiss` on any drift, naming the recorded and the actual call; a passed test must also consume its whole recording, or teardown fails it naming the first leftover interaction. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy Deliberately not here yet: canonical content-based match keys (LIT-5741), streaming chunk fidelity (LIT-5742), and scoping record/replay to provider-bound traffic (LIT-5745) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 6b27bb459a5..da2a7da0bfa 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -15,7 +15,7 @@ shared fixtures build on it. import functools import os -from collections.abc import Iterator +from collections.abc import Generator, Iterator from datetime import datetime, timezone import pytest @@ -27,6 +27,7 @@ from fixture_transport import ( fixture_mode_collection_error, fixture_report_lines, parse_fixture_mode, + replay_leftover_error, ) from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager @@ -34,6 +35,7 @@ from proxy_client import ProxyClient, build_proxy_client _E2E_TEST_RAN = pytest.StashKey[bool]() +_CALL_PASSED = pytest.StashKey[bool]() def pytest_configure(config: pytest.Config) -> None: @@ -134,6 +136,36 @@ def pytest_runtest_call(item: pytest.Item) -> None: item.session.stash[_E2E_TEST_RAN] = True +@pytest.hookimpl(wrapper=True) +def pytest_runtest_makereport( + item: pytest.Item, call: pytest.CallInfo[None] +) -> Generator[None, pytest.TestReport, pytest.TestReport]: + """Stash the call-phase outcome so teardown can tell a passed test from a + failed one without re-deriving it.""" + report = yield + if report.when == "call": + item.stash[_CALL_PASSED] = report.passed + return report + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_teardown(item: pytest.Item) -> Generator[None, None, None]: + """In replay mode a passing test must consume its whole recording: leftover + interactions mean the test now makes fewer calls than it did at record time, + so the replay proved less than the bundle claims. The check runs after the + yield so fixture finalizers replay their recorded calls first. Failed tests + are left alone - their own failure already explains any unconsumed tail.""" + result = yield + if not item.stash.get(_CALL_PASSED, False): + return result + reason = replay_leftover_error( + mode_raw=FIXTURE_MODE_RAW, bundle_dir=FIXTURE_DIR, test_key=item.nodeid + ) + if reason is not None: + pytest.fail(reason) + return result + + def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: """Once the whole e2e session is done (all suites), optionally truncate the spend logs so the DB doesn't accumulate test rows. The truncate is destructive diff --git a/tests/e2e/fixture_transport.py b/tests/e2e/fixture_transport.py index 756362cf29e..b99b0d2d80d 100644 --- a/tests/e2e/fixture_transport.py +++ b/tests/e2e/fixture_transport.py @@ -332,6 +332,21 @@ class ReplaySource: self._cursors[slug] = index + 1 return interaction + def leftover_error(self, test_key: str) -> str | None: + """Non-None when the test consumed fewer interactions than were recorded, + meaning a passing replay proved less than the bundle claims.""" + slug = slug_for_test(test_key) + recorded = self.bundle.interactions.get(slug, ()) + consumed = self._cursors.get(slug, 0) + if consumed >= len(recorded): + return None + pending = recorded[consumed] + return ( + f"replay incomplete for {test_key}: {len(recorded) - consumed} of {len(recorded)} recorded " + f"interactions never consumed, next is {pending.request.method} {pending.request.path}; " + "re-record with E2E_FIXTURE_MODE=record" + ) + def _expect_result(interaction: Interaction) -> RecordedResult: match interaction.response: @@ -474,6 +489,15 @@ def _shared_replay_source(root: Path) -> ReplaySource: return ReplaySource(bundle=loaded) +def replay_leftover_error(*, mode_raw: str, bundle_dir: Path, test_key: str) -> str | None: + """Teardown-time completeness check: in replay mode a passed test with + unconsumed recorded interactions must fail instead of passing against a + recording it no longer matches. Inert in every other mode.""" + if parse_fixture_mode(mode_raw) != "replay": + return None + return _shared_replay_source(bundle_dir).leftover_error(test_key) + + def select_transport( live: Transport, *, mode_raw: str, bundle_dir: Path, master_key: str ) -> Transport: diff --git a/tests/e2e/test_fixture_transport.py b/tests/e2e/test_fixture_transport.py index 6ffcdaeb95f..5c7201cca37 100644 --- a/tests/e2e/test_fixture_transport.py +++ b/tests/e2e/test_fixture_transport.py @@ -50,6 +50,7 @@ from fixture_transport import ( fixture_mode_collection_error, fixture_report_lines, parse_fixture_mode, + replay_leftover_error, select_transport, ) from transport import Transport @@ -354,6 +355,52 @@ class TestReplayTransport: replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) +class TestReplayLeftover: + def test_fully_consumed_recording_leaves_nothing(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + source = replay_source(root) + replay: Transport = ReplayTransport(source=source, master_key="sk-1234") + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + assert source.leftover_error(current_test_key()) is None + + def test_unconsumed_trailing_interactions_name_the_next_call(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + recording.probe("/health/liveliness", params=Query(q="1")) + source = replay_source(root) + replay: Transport = ReplayTransport(source=source, master_key="sk-1234") + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + error = source.leftover_error(current_test_key()) + assert error is not None + assert "1 of 2 recorded interactions never consumed" in error + assert "next is probe /health/liveliness" in error + assert "re-record with E2E_FIXTURE_MODE=record" in error + + def test_test_without_recordings_has_no_leftover(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + make_recorder(root) + assert replay_source(root).leftover_error("suite.py::test_never_recorded") is None + + def test_inert_outside_replay_mode(self, tmp_path: Path) -> None: + missing = tmp_path / "missing" + assert replay_leftover_error(mode_raw="", bundle_dir=missing, test_key="k") is None + assert replay_leftover_error(mode_raw="record", bundle_dir=missing, test_key="k") is None + + def test_replay_mode_reads_the_shared_bundle(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + error = replay_leftover_error(mode_raw="replay", bundle_dir=root, test_key=current_test_key()) + assert error is not None + assert "1 of 1 recorded interactions never consumed" in error + + class TestSelectTransport: def test_live_returns_the_live_transport_untouched(self, tmp_path: Path) -> None: fake = FakeTransport() From 17b72d5089c7ba13c23e45837fd06a28825a82e8 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 18 Aug 2026 23:26:41 +0000 Subject: [PATCH 20/88] fix(search): send MCP-Protocol-Version on AgentCore gateway calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 9 ++++- .../test_agentcore_search_transformation.py | 39 ++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index ca9759ed151..5567dc8403e 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -66,6 +66,11 @@ AGENTCORE_DEFAULT_TOOL_NAME: Final = "web-search-tool___WebSearch" # with the proxy's credentials. AGENTCORE_TOOL_NAME_SUFFIX: Final = "___WebSearch" +# MCP revision this provider speaks. Sent on every request because the gateway is +# called statelessly, without an initialize handshake to negotiate a version; +# servers that predate the header ignore it. +AGENTCORE_MCP_PROTOCOL_VERSION: Final = "2025-06-18" + _GATEWAY_REGION_PATTERN: Final = re.compile(r"\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com") _SSE_EVENT_SEPARATOR: Final = re.compile(r"\n[ \t]*\n") @@ -147,7 +152,8 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): ) -> dict: # mutable-ok: the handler passes these headers straight to httpx, which wants a dict """ Set MCP transport headers. Per the MCP Streamable HTTP transport spec, - the client MUST accept both application/json and text/event-stream. + the client MUST accept both application/json and text/event-stream, and + declare its protocol revision with MCP-Protocol-Version. Authentication itself happens in sign_request(): bearer token for CUSTOM_JWT gateways, AWS SigV4 for AWS_IAM gateways. @@ -156,6 +162,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): **headers, "Content-Type": "application/json", "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": AGENTCORE_MCP_PROTOCOL_VERSION, } def get_complete_url( diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 38189abe0e8..63ec2286c3d 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -13,7 +13,10 @@ import pytest from unittest.mock import AsyncMock, patch, MagicMock import litellm -from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig +from litellm.llms.bedrock.search.transformation import ( + AGENTCORE_MCP_PROTOCOL_VERSION, + AgentCoreSearchConfig, +) GATEWAY_URL = "https://testgateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" @@ -148,11 +151,43 @@ class TestAgentCoreSearch: assert config.get_complete_url(api_base=GATEWAY_URL, optional_params={}) == GATEWAY_URL def test_validate_environment_sets_mcp_headers(self): - """MCP Streamable HTTP requires accepting both JSON and SSE.""" + """MCP Streamable HTTP requires accepting both JSON and SSE, and declaring + the protocol revision the client speaks.""" config = AgentCoreSearchConfig() headers = config.validate_environment(headers={}) assert headers["Accept"] == "application/json, text/event-stream" assert headers["Content-Type"] == "application/json" + assert headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION + + def test_protocol_version_header_survives_signing(self): + """Both auth paths must keep the MCP-Protocol-Version header on the wire.""" + config = AgentCoreSearchConfig() + headers = config.validate_environment(headers={}) + + bearer_headers, _ = config.sign_request( + headers=headers, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + api_key="test-jwt-token", + ) + assert bearer_headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION + + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "AKIAIOSFODNN7EXAMPLE", + "AWS_SECRET_ACCESS_KEY": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + }, + ): + signed_headers, _ = config.sign_request( + headers=headers, + optional_params={"aws_region_name": "us-east-1"}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert signed_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert signed_headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION def test_transform_search_response_parses_sse_frame(self): """Gateway may answer with an SSE-framed JSON-RPC message.""" From 608d7499836c1aaa7fe752c473c68d5813b9f29b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:46:59 -0700 Subject: [PATCH 21/88] fix(batches): stop one bad output line from zeroing an entire batch's spend --- litellm/batches/batch_utils.py | 161 ++++++++++++------ .../test_litellm/batches/test_batch_utils.py | 40 +++-- .../proxy/hooks/test_batch_file_validation.py | 12 +- type-discipline-budget.json | 2 +- 4 files changed, 149 insertions(+), 66 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index c2cbb9604e5..feb84ccd8a6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from typing import Any, Final, Literal @@ -87,7 +87,7 @@ async def _handle_completed_batch( return batch_cost, batch_usage, [model_name] return _aggregate_batch_cost_usage_models( - entries=_iter_batch_input_entries(file_content), + entries=_iter_batch_output_entries(file_content), custom_llm_provider=custom_llm_provider, model_name=model_name, model_info=model_info, @@ -111,43 +111,91 @@ def _iter_successful_output_line_stats( model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats]: + for entry in entries: + stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) + if stats is not None: + yield stats + + +def _safe_output_line_stats( + entry: Mapping, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + model_name: str | None, + model_info: ModelInfo | None, +) -> _BatchOutputLineStats | None: + """Return the stats for one batch output line, or None for a line that is + unsuccessful or cannot be costed, so a single bad line never aborts the + whole batch's cost accounting.""" + custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None + try: + if not _batch_response_was_successful(entry, custom_llm_provider): + return None + return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info) + except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch + verbose_logger.warning( + "batch output line could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed. custom_id=%s error=%s", + custom_id, + str(e), + ) + return None + + +def _compute_output_line_stats( + entry: Mapping, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + model_name: str | None, + model_info: ModelInfo | None, +) -> _BatchOutputLineStats: + response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) + usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) + prompt_details: Final = parse_prompt_tokens_details(usage) + raw_model: Final = response_body.get("model") + response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None + return _BatchOutputLineStats( + cost=_output_line_cost( + response_body=response_body, + usage=usage, + custom_llm_provider=custom_llm_provider, + model_name=model_name, + response_model=response_model, + model_info=model_info, + ), + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens, + cache_read_tokens=prompt_details["cache_hit_tokens"], + cache_creation_tokens=prompt_details["cache_creation_tokens"], + model=response_model, + ) + + +def _output_line_cost( + response_body: Mapping, + usage: Usage, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + model_name: str | None, + response_model: str | None, + model_info: ModelInfo | None, +) -> float: from litellm.cost_calculator import batch_cost_calculator - for entry in entries: - if not _batch_response_was_successful(entry, custom_llm_provider): - continue - response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider) - usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) - prompt_details = parse_prompt_tokens_details(usage) - raw_model = response_body.get("model") - response_model = raw_model if isinstance(raw_model, str) and raw_model else None - if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"): - if custom_llm_provider == "bedrock" and model_name: - cost_model = model_name - else: - cost_model = response_model or model_name or "" - prompt_cost, completion_cost = batch_cost_calculator( - usage=usage, - model=cost_model, - custom_llm_provider=custom_llm_provider, - model_info=model_info, - ) - line_cost = prompt_cost + completion_cost - else: - line_cost = litellm.completion_cost( - completion_response=response_body, - custom_llm_provider=custom_llm_provider, - call_type=CallTypes.aretrieve_batch.value, - ) - yield _BatchOutputLineStats( - cost=line_cost, - prompt_tokens=usage.prompt_tokens, - completion_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, - cache_read_tokens=prompt_details["cache_hit_tokens"], - cache_creation_tokens=prompt_details["cache_creation_tokens"], - model=response_model, + if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"): + return litellm.completion_cost( + completion_response=response_body, + custom_llm_provider=custom_llm_provider, + call_type=CallTypes.aretrieve_batch.value, ) + cost_model: Final = ( + model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" + ) + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, + model=cost_model, + custom_llm_provider=custom_llm_provider, + model_info=model_info, + ) + return prompt_cost + completion_cost def _aggregate_batch_cost_usage_models( @@ -338,9 +386,10 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]: """ - Get the file content as a list of dictionaries from JSON Lines format + Get the file content as a list of dictionaries from JSON Lines format, + skipping malformed lines """ - return list(_iter_batch_input_entries(file_content)) + return list(_iter_batch_output_entries(file_content)) def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]: @@ -361,15 +410,29 @@ def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]: yield line -def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]: +def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]: """ - Yield parsed batch input JSONL entries one at a time without materializing the - whole file as a list, so peak memory stays bounded. Raises on a malformed line; - callers that must survive bad rows should iterate ``_iter_batch_input_lines`` - and parse per-row instead. + Yield parsed batch output JSONL entries one at a time without materializing + the whole file as a list, so peak memory stays bounded. A malformed or + non-object line is skipped with a warning so one bad line never aborts the + whole batch's cost accounting. """ for line in _iter_batch_input_lines(file_content): - yield json.loads(line) + entry = _parse_batch_output_line(line) + if entry is not None: + yield entry + + +def _parse_batch_output_line(line: bytes) -> dict | None: + try: + parsed: Final = json.loads(line) + except json.JSONDecodeError as e: + verbose_logger.warning("skipping malformed batch output line: %s", str(e)) + return None + if isinstance(parsed, dict): + return parsed + verbose_logger.warning("skipping non-object batch output line of type %s", type(parsed).__name__) + return None # A batch request's input tokens scale roughly with its serialized size, so this @@ -440,7 +503,7 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int: return 0 -def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_provider: str = "openai") -> Usage: +def _get_batch_job_usage_from_response_body(response_body: Mapping, custom_llm_provider: str = "openai") -> Usage: """ Get the tokens of a batch job from the response body """ @@ -472,7 +535,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov return usage -def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> dict: +def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping) -> dict: """ Get the ``result`` object from a line of an Anthropic message batch results JSONL file. @@ -482,7 +545,9 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> d return batch_results_line.get("result", None) or {} -def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> Any: +def _get_response_from_batch_job_output_file( + batch_job_output_file: Mapping, custom_llm_provider: str = "openai" +) -> Any: """ Get the response from the batch job output file """ @@ -495,7 +560,7 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom return _response_body -def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> bool: +def _batch_response_was_successful(batch_job_output_file: Mapping, custom_llm_provider: str = "openai") -> bool: """ Check if the batch job response was successful diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 573882ebfca..254d663af93 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -150,13 +150,13 @@ def test_parse_jsonl_empty_content_is_empty_list(): assert bu._get_file_content_as_dictionary(b"") == [] -def test_parse_jsonl_malformed_raises(): - with pytest.raises(Exception): - bu._get_file_content_as_dictionary(b"not valid json") +def test_parse_jsonl_malformed_lines_skipped(): + content = b'{"a": 1}\nnot valid json\n{"b": 2}\n' + assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}] # =========================================================================== # -# _iter_batch_input_lines / _iter_batch_input_entries (JSONL parsing) +# _iter_batch_input_lines / _iter_batch_output_entries (JSONL parsing) # =========================================================================== # @@ -173,19 +173,17 @@ def test_iter_input_lines_empty(): assert list(bu._iter_batch_input_lines(b"")) == [] -def test_iter_input_entries_parses_each_row(): +def test_iter_output_entries_parses_each_row(): content = b'{"body": {"model": "gpt-4o"}}\n{"body": {"model": "claude-3"}}\n' - assert list(bu._iter_batch_input_entries(content)) == [ + assert list(bu._iter_batch_output_entries(content)) == [ {"body": {"model": "gpt-4o"}}, {"body": {"model": "claude-3"}}, ] -def test_iter_input_entries_raises_on_malformed_line(): - # _iter_batch_input_entries raises on a bad row; callers that must survive - # bad rows iterate _iter_batch_input_lines and parse per-row instead. - with pytest.raises(Exception): - list(bu._iter_batch_input_entries(b'{"ok":1}\nnot-json\n')) +def test_iter_output_entries_skips_malformed_and_non_object_lines(): + content = b'{"ok": 1}\nnot-json\n[1, 2]\n{"ok": 2}\n' + assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}] # =========================================================================== # @@ -471,6 +469,26 @@ def test_cost_from_content_completion_cost_path(monkeypatch): assert len(calls) == 2 # failed row not costed +def test_empty_body_line_does_not_zero_whole_batch(): + # Regression: a status-200 row with an empty body made the real + # litellm.completion_cost raise ValueError, aborting the aggregation so the + # entire batch was booked at $0. The bad line must be skipped instead. + rows = [ + _success_row(usage=_usage(10, 5)), + { + "custom_id": "request-poison-empty", + "response": {"status_code": 200, "request_id": "inject-empty-body", "body": {}}, + }, + _success_row(usage=_usage(20, 10)), + ] + + cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + + assert cost > 0.0 + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) + assert models == ["gpt-4o", "gpt-4o"] + + def test_cost_from_content_model_info_path(monkeypatch): # model_info set -> batch_cost_calculator(prompt_cost, completion_cost). import litellm.cost_calculator as cc diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 1ce1a2f3e51..28624254565 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -1739,18 +1739,18 @@ def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes: return ("\n".join(rows)).encode("utf-8") -def test_iter_batch_input_entries_matches_dict_list(): +def test_iter_batch_output_entries_matches_dict_list(): from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, - _iter_batch_input_entries, + _iter_batch_output_entries, ) raw = _make_batch_input_bytes(50) - streamed = list(_iter_batch_input_entries(raw)) + streamed = list(_iter_batch_output_entries(raw)) assert streamed == _get_file_content_as_dictionary(raw) assert streamed[0]["custom_id"] == "request-0" # tolerant of blank lines and a missing trailing newline - assert list(_iter_batch_input_entries(raw + b"\n\n")) == streamed + assert list(_iter_batch_output_entries(raw + b"\n\n")) == streamed def test_streaming_count_peak_below_dict_list(): @@ -1759,7 +1759,7 @@ def test_streaming_count_peak_below_dict_list(): from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, - _iter_batch_input_entries, + _iter_batch_output_entries, ) raw = _make_batch_input_bytes(8000) @@ -1777,7 +1777,7 @@ def test_streaming_count_peak_below_dict_list(): def _stream(): count = 0 models: set = set() - for entry in _iter_batch_input_entries(raw): + for entry in _iter_batch_output_entries(raw): count += 1 model = (entry.get("body") or {}).get("model") if model: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f8e481dc142..43753224714 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22894 + "limit": 22891 }, "LIT002": { "limit": 26888 From f614f039c56c753c3ca0ad8c884eaa81e5976da3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:00:47 -0700 Subject: [PATCH 22/88] fix(ui): drop stale user search answers so Enter commits the current match The Add Member modal and the Create Key owner picker both search users server-side on a 300ms debounce with nothing sequencing the requests, so a slow answer to an earlier, shorter search can land after the current one and replace the list. With the first row highlighted at all times, Enter then commits whoever sits on top of that abandoned batch, ordered newest account first rather than best match. Each search now takes a sequence number and only the newest one is allowed to reach the option list or clear the spinner. --- .../user_search_modal.test.tsx | 55 +++++++++++++++++++ .../common_components/user_search_modal.tsx | 10 +++- .../create_key_button.integration.test.tsx | 34 +++++++++++- .../organisms/create_key_button.tsx | 10 +++- 4 files changed, 104 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx index cbaf00b7131..98b85c738b4 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -198,3 +198,58 @@ describe("UserSearchModal submit payload", () => { await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); }); }); + +describe("UserSearchModal out-of-order search results", () => { + const answers = new Map void>(); + + beforeEach(() => { + answers.clear(); + vi.mocked(userFilterUICall).mockReset(); + vi.mocked(userFilterUICall).mockImplementation( + (_accessToken, params) => + new Promise((resolve) => { + answers.set(params.get("user_email") ?? "", resolve); + }) as never, + ); + }); + + const answerFor = async (search: string, users: { user_id: string; user_email: string }[]) => { + const resolve = answers.get(search); + if (resolve === undefined) throw new Error(`no pending search for "${search}"`); + await act(async () => { + resolve(users); + }); + }; + + it("commits the current search's match when an abandoned search answers last", async () => { + const user = userEvent.setup(); + const onSubmit = vi.fn(); + render(); + + const input = getEmailSearchInput(); + await user.click(input); + await user.type(input, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + + await user.type(input, "ce.smith@example.com"); + await waitFor(() => expect(answers.has("alice.smith@example.com")).toBe(true), { timeout: 3000 }); + + await answerFor("alice.smith@example.com", [{ user_id: "u-smith", user_email: "alice.smith@example.com" }]); + await screen.findByRole("option", { name: "alice.smith@example.com" }); + + await answerFor("ali", [ + { user_id: "u-jones", user_email: "alice.jones@example.com" }, + { user_id: "u-smith", user_email: "alice.smith@example.com" }, + ]); + + await user.keyboard("{Enter}"); + await user.click(screen.getByRole("button", { name: /add member/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + expect(onSubmit.mock.calls[0][0]).toStrictEqual({ + user_email: "alice.smith@example.com", + user_id: "u-smith", + role: "user", + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index aeead9d82e6..71d38c6db25 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -1,4 +1,4 @@ -import { useState } from "react"; +import { useRef, useState } from "react"; import { Modal, Alert } from "antd"; import { UserAddOutlined } from "@ant-design/icons"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; @@ -78,8 +78,13 @@ const UserSearchModal: React.FC = ({ const [loading, setLoading] = useState(false); const [selectedField, setSelectedField] = useState<"user_email" | "user_id">("user_email"); const [isSubmitting, setIsSubmitting] = useState(false); + const latestSearchRef = useRef(0); const fetchUsers = async (searchText: string, fieldName: "user_email" | "user_id"): Promise => { + const searchId = latestSearchRef.current + 1; + latestSearchRef.current = searchId; + const isLatestSearch = (): boolean => searchId === latestSearchRef.current; + if (!searchText) { setUserOptions([]); return; @@ -96,6 +101,7 @@ const UserSearchModal: React.FC = ({ return; } const response = await userFilterUICall(accessToken, params); + if (!isLatestSearch()) return; const data: User[] = response; const options: UserOption[] = data.map((user) => ({ @@ -107,7 +113,7 @@ const UserSearchModal: React.FC = ({ } catch (error) { console.error("Error fetching users:", error); } finally { - setLoading(false); + if (isLatestSearch()) setLoading(false); } }; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 9788d365fd0..c413bbeebfc 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -1,6 +1,6 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import { act, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import type { Team } from "../key_team_helpers/key_list"; import { keyCreateCall, keyCreateServiceAccountCall, modelAvailableCall, userFilterUICall } from "../networking"; import CreateKey from "./create_key_button"; @@ -803,6 +803,38 @@ describe("CreateKey", () => { vi.useRealTimers(); } }); + + it("keeps the current search's users when an abandoned search answers last", async () => { + const answers = new Map void>(); + vi.mocked(userFilterUICall).mockImplementation( + (_accessToken, params) => + new Promise((resolve) => { + answers.set(params.get("user_email") ?? "", resolve); + }) as never, + ); + + const user = userEvent.setup(); + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + const search = antdSearchInput(await screen.findByText("Type email to search for users")); + + await user.type(search, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + + await user.type(search, "ce.smith@example.com"); + await waitFor(() => expect(answers.has("alice.smith@example.com")).toBe(true), { timeout: 3000 }); + + await act(async () => { + answers.get("alice.smith@example.com")?.([{ user_id: "u-smith", user_email: "alice.smith@example.com" }]); + }); + await screen.findByTitle("alice.smith@example.com (u-smith)"); + + await act(async () => { + answers.get("ali")?.([{ user_id: "u-jones", user_email: "alice.jones@example.com" }]); + }); + + expect(screen.queryByTitle("alice.jones@example.com (u-jones)")).not.toBeInTheDocument(); + expect(screen.getByTitle("alice.smith@example.com (u-smith)")).toBeInTheDocument(); + }); }); describe("created key display", () => { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 77ef31e0ede..1e7fff687f1 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -27,7 +27,7 @@ import { import { ChevronDown } from "lucide-react"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; -import React, { useEffect, useState } from "react"; +import React, { useEffect, useRef, useState } from "react"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; @@ -195,6 +195,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const [possibleUIRoles, setPossibleUIRoles] = useState>>({}); const [userOptions, setUserOptions] = useState([]); const [userSearchLoading, setUserSearchLoading] = useState(false); + const latestUserSearchRef = useRef(0); const [disabledCallbacks, setDisabledCallbacks] = useState([]); const [keyType, setKeyType] = useState("llm_api"); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); @@ -503,6 +504,10 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const fetchUsers = async (searchText: string): Promise => { + const searchId = latestUserSearchRef.current + 1; + latestUserSearchRef.current = searchId; + const isLatestSearch = (): boolean => searchId === latestUserSearchRef.current; + if (!searchText) { setUserOptions([]); return; @@ -516,6 +521,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp return; } const response = await userFilterUICall(accessToken, params); + if (!isLatestSearch()) return; const data: User[] = response; const options: UserOption[] = data.map((user) => ({ @@ -529,7 +535,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp console.error("Error fetching users:", error); toast.fromError("Failed to search for users"); } finally { - setUserSearchLoading(false); + if (isLatestSearch()) setUserSearchLoading(false); } }; From 5c6391d7d2590f19e6cd5bc2f5d89c5aa4ffe5f1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:15:22 -0700 Subject: [PATCH 23/88] refactor(batches): parameterize the batch output Mapping annotations --- litellm/batches/batch_utils.py | 18 +++++++++++------- tests/test_litellm/batches/test_batch_utils.py | 5 ++--- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index feb84ccd8a6..0f8acce3379 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -118,7 +118,7 @@ def _iter_successful_output_line_stats( def _safe_output_line_stats( - entry: Mapping, + entry: Mapping[str, Any], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, @@ -142,7 +142,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( - entry: Mapping, + entry: Mapping[str, Any], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, @@ -171,7 +171,7 @@ def _compute_output_line_stats( def _output_line_cost( - response_body: Mapping, + response_body: Mapping[str, Any], usage: Usage, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, @@ -503,7 +503,9 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int: return 0 -def _get_batch_job_usage_from_response_body(response_body: Mapping, custom_llm_provider: str = "openai") -> Usage: +def _get_batch_job_usage_from_response_body( + response_body: Mapping[str, Any], custom_llm_provider: str = "openai" +) -> Usage: """ Get the tokens of a batch job from the response body """ @@ -535,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: Mapping, custom_llm_p return usage -def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping) -> dict: +def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict: """ Get the ``result`` object from a line of an Anthropic message batch results JSONL file. @@ -546,7 +548,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping) - def _get_response_from_batch_job_output_file( - batch_job_output_file: Mapping, custom_llm_provider: str = "openai" + batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai" ) -> Any: """ Get the response from the batch job output file @@ -560,7 +562,9 @@ def _get_response_from_batch_job_output_file( return _response_body -def _batch_response_was_successful(batch_job_output_file: Mapping, custom_llm_provider: str = "openai") -> bool: +def _batch_response_was_successful( + batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai" +) -> bool: """ Check if the batch job response was successful diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 254d663af93..a30420e6224 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -470,9 +470,8 @@ def test_cost_from_content_completion_cost_path(monkeypatch): def test_empty_body_line_does_not_zero_whole_batch(): - # Regression: a status-200 row with an empty body made the real - # litellm.completion_cost raise ValueError, aborting the aggregation so the - # entire batch was booked at $0. The bad line must be skipped instead. + """A status-200 row with an empty body makes litellm.completion_cost raise; + that line must be skipped instead of zeroing the whole batch.""" rows = [ _success_row(usage=_usage(10, 5)), { From 5a0a8ffafe6b8af4c73cc28b5cbb17376a835350 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:17:15 +0000 Subject: [PATCH 24/88] fix(model_prices): correct gemini and deepseek pricing and add deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 54 ++++++++++--------- model_prices_and_context_window.json | 54 ++++++++++--------- 2 files changed, 60 insertions(+), 48 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f9027313b..1a1652451a1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10004,6 +10004,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -15042,6 +15043,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -20562,8 +20564,8 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-image": { - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -20571,8 +20573,8 @@ "mode": "image_generation", "output_cost_per_image": 0.045, "output_cost_per_image_token": 6e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "rpm": 1000, "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", @@ -20605,8 +20607,8 @@ }, "gemini/gemini-3.1-flash-image-preview": { "deprecation_date": "2026-06-25", - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -20614,8 +20616,8 @@ "mode": "image_generation", "output_cost_per_image": 0.045, "output_cost_per_image_token": 6e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "rpm": 1000, "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image-preview", @@ -23245,6 +23247,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -36161,6 +36164,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -36170,6 +36174,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -36179,6 +36184,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -47975,15 +47981,15 @@ }, "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 2.8e-09, - "input_cost_per_token": 1.4e-07, - "input_cost_per_token_cache_hit": 2.8e-09, + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 1.32e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48001,15 +48007,15 @@ }, "deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 8.7e-07, + "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48027,15 +48033,15 @@ }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 2.8e-09, - "input_cost_per_token": 1.4e-07, - "input_cost_per_token_cache_hit": 2.8e-09, + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 1.32e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48053,15 +48059,15 @@ }, "deepseek/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 8.7e-07, + "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 07f9027313b..1a1652451a1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10004,6 +10004,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -15042,6 +15043,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -20562,8 +20564,8 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-image": { - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -20571,8 +20573,8 @@ "mode": "image_generation", "output_cost_per_image": 0.045, "output_cost_per_image_token": 6e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "rpm": 1000, "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", @@ -20605,8 +20607,8 @@ }, "gemini/gemini-3.1-flash-image-preview": { "deprecation_date": "2026-06-25", - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -20614,8 +20616,8 @@ "mode": "image_generation", "output_cost_per_image": 0.045, "output_cost_per_image_token": 6e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "rpm": 1000, "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image-preview", @@ -23245,6 +23247,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -36161,6 +36164,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -36170,6 +36174,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -36179,6 +36184,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -47975,15 +47981,15 @@ }, "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 2.8e-09, - "input_cost_per_token": 1.4e-07, - "input_cost_per_token_cache_hit": 2.8e-09, + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 1.32e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48001,15 +48007,15 @@ }, "deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 8.7e-07, + "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48027,15 +48033,15 @@ }, "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 2.8e-09, - "input_cost_per_token": 1.4e-07, - "input_cost_per_token_cache_hit": 2.8e-09, + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "input_cost_per_token_cache_hit": 1.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 1.32e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -48053,15 +48059,15 @@ }, "deepseek/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 8.7e-07, + "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" From dcd8bb3f38561c0030499a0fa074cddee1aed4e8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:47:44 +0000 Subject: [PATCH 25/88] test(model_prices): update DeepSeek V4 pricing expectations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index efccdc4a986..0218f8f6219 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -3787,8 +3787,8 @@ def test_deepseek_v4_models_in_cost_map(): configured in model_prices_and_context_window.json. Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - - deepseek-v4-flash: $0.14/M input, $0.28/M output - - deepseek-v4-pro: $0.435/M input, $0.87/M output (75% discounted active price) + - deepseek-v4-flash: $0.44/M input, $1.32/M output + - deepseek-v4-pro: $1.32/M input, $3.96/M output Closes https://github.com/BerriAI/litellm/issues/26709 """ @@ -3801,8 +3801,8 @@ def test_deepseek_v4_models_in_cost_map(): # --- bare model names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09), - ("deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09), + ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from model_prices_and_context_window.json" @@ -3817,8 +3817,8 @@ def test_deepseek_v4_models_in_cost_map(): # --- provider-prefixed names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09), - ("deepseek/deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09), + ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from model_prices_and_context_window.json" @@ -3845,8 +3845,8 @@ def test_deepseek_v4_models_in_backup_cost_map(): # --- bare model names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09), - ("deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09), + ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from backup JSON" @@ -3859,8 +3859,8 @@ def test_deepseek_v4_models_in_backup_cost_map(): # --- provider-prefixed names --- for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09), - ("deepseek/deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09), + ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), + ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from backup JSON" From ae18f055ee3ce5983b3023d533da7b7cc9676d22 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 19 Aug 2026 19:07:17 +0000 Subject: [PATCH 26/88] fix(search): harden AgentCore gateway trust, error and SSE handling Refuse to SigV4-sign requests to hosts that are neither an AgentCore gateway hostname nor AGENTCORE_GATEWAY_URL's host, match gateway hostnames on the URL host instead of anywhere in the URL, accept the env token when api_base is a real gateway, raise on tools/call responses with result.isError, and split CRLF-framed SSE events. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 52 ++++++- .../test_agentcore_search_transformation.py | 141 +++++++++++++++--- 2 files changed, 161 insertions(+), 32 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 5567dc8403e..85c5818031b 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -71,13 +71,19 @@ AGENTCORE_TOOL_NAME_SUFFIX: Final = "___WebSearch" # servers that predate the header ignore it. AGENTCORE_MCP_PROTOCOL_VERSION: Final = "2025-06-18" -_GATEWAY_REGION_PATTERN: Final = re.compile(r"\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com") +# Matched against the URL host so a crafted path or query string can't pass for +# a gateway hostname. +_GATEWAY_HOST_PATTERN: Final = re.compile(r"[a-z0-9-]+\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com") -_SSE_EVENT_SEPARATOR: Final = re.compile(r"\n[ \t]*\n") +_SSE_EVENT_SEPARATOR: Final = re.compile(r"\r?\n[ \t]*\r?\n") _SSE_LINE_PREFIXES: Final = ("event:", "data:", ":", "id:", "retry:") +def _gateway_host_match(api_base: str) -> re.Match[str] | None: + return _GATEWAY_HOST_PATTERN.fullmatch(httpx.URL(api_base).host) + + def _string_field(item: Mapping[str, object], *keys: str) -> str | None: return next( (value for key in keys if isinstance(value := item.get(key), str) and value), @@ -245,16 +251,17 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): if not isinstance(request_data, dict): raise TypeError("AgentCore search expects a single dict request body") - # Server-managed token fallback is gated on the request targeting the - # operator-configured gateway host, otherwise an authenticated caller - # could point api_base at their own server (e.g. via - # /search_tools/test_connection) and receive AGENTCORE_GATEWAY_TOKEN. + # Server-managed credentials only go to a trusted host, otherwise an + # authenticated caller could point api_base at their own server (e.g. via + # /search_tools/test_connection) and collect AGENTCORE_GATEWAY_TOKEN or a + # SigV4 signature with the proxy's credential scope and session token. + gateway_host_match: Final = _gateway_host_match(api_base) bearer_token: Final = self.resolve_server_api_key( caller_api_key=api_key, caller_api_base=api_base, key_env_vars=("AGENTCORE_GATEWAY_TOKEN",), base_env_var="AGENTCORE_GATEWAY_URL", - default_api_base=None, + default_api_base=api_base if gateway_host_match else None, ) if bearer_token: bearer_headers: Final = { # mutable-ok: httpx request headers are a dict @@ -263,6 +270,13 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): } return bearer_headers, json.dumps(request_data).encode() + if gateway_host_match is None and not self._is_configured_gateway(api_base): + raise ValueError( + f"Refusing to send SigV4-signed AgentCore requests to '{api_base}': it is neither an " + "AgentCore gateway hostname nor the host in AGENTCORE_GATEWAY_URL. Set " + "AGENTCORE_GATEWAY_URL to authorize a custom gateway hostname." + ) + signing_params: Final = ( optional_params if optional_params.get("aws_region_name") is not None @@ -284,6 +298,13 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): api_key="", ) + @staticmethod + def _is_configured_gateway(api_base: str) -> bool: + configured: Final = get_secret_str("AGENTCORE_GATEWAY_URL") + if not configured: + return False + return httpx.URL(configured).host == httpx.URL(api_base).host + @staticmethod def _signing_region(api_base: str) -> str: """ @@ -296,7 +317,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): nothing rather than silently signing for a guessed region the gateway would reject with a confusing auth error. """ - match: Final = _GATEWAY_REGION_PATTERN.search(api_base) + match: Final = _gateway_host_match(api_base) if match: return match.group(1) @@ -336,6 +357,15 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): message=f"AgentCore gateway MCP error: {error}", ) + # A failed tools/call is reported in-band, as HTTP 200 with result.isError + # and the failure text where the results would be. + result: Final = response_json.get("result") + if isinstance(result, dict) and result.get("isError"): + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", + ) + return SearchResponse( results=[ # mutable-ok: SearchResponse.results is a pydantic list field _to_search_result(item) @@ -345,6 +375,12 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): object="search", ) + def _tool_error_message(self, response_json: Mapping[str, object]) -> str: + texts: Final = tuple( + text for block in self._text_blocks(response_json) if isinstance(text := block.get("text"), str) + ) + return " ".join(texts) if texts else json.dumps(response_json.get("result"))[:500] + @staticmethod def _text_blocks(response_json: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: result: Final = response_json.get("result") diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 63ec2286c3d..889f78c58b9 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -231,6 +231,37 @@ class TestAgentCoreSearch: with pytest.raises(Exception, match="tool not found"): config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + def test_transform_search_response_raises_on_tool_error(self): + """A failed tools/call comes back as HTTP 200 with result.isError; it must not be + reported to the caller as a successful search with zero results.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "isError": True, + "content": [{"type": "text", "text": "AccessDeniedException: not authorized"}], + }, + } + ) + with pytest.raises(Exception, match="AccessDeniedException"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_transform_search_response_parses_crlf_framed_sse(self): + """SSE streams may be CRLF framed; events must still split into separate events.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\r\ndata: {json.dumps(progress)}\r\n\r\n" + f"event: message\r\ndata: {json.dumps(_mcp_response_body())}\r\n\r\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + def test_sign_request_uses_bearer_token_when_api_key_set(self): """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" config = AgentCoreSearchConfig() @@ -280,6 +311,54 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) + def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self): + """api_base pointing at a real gateway is a trusted destination for the env token, + so operators configuring api_base in yaml don't also need AGENTCORE_GATEWAY_URL.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + + @pytest.mark.parametrize( + "untrusted_api_base", + [ + "https://attacker.example.com/mcp", + # gateway hostname in the path/query must not pass for the host + "https://attacker.example.com/gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", + ], + ) + def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base): + """A SigV4 signature carries the proxy's credential scope and session token, so it + must never be sent to a host that is not the operator's gateway.""" + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-1"}, + request_data={"jsonrpc": "2.0"}, + api_base=untrusted_api_base, + ) + mock_base_sign.assert_not_called() + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + def test_sign_request_does_not_leak_bedrock_bearer_token(self): """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not replace SigV4 on requests to an AgentCore gateway.""" @@ -303,39 +382,49 @@ class TestAgentCoreSearch: def test_sign_request_custom_hostname_requires_region(self): """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() + custom_url = "https://gateway.internal.example.com/mcp" + os.environ["AGENTCORE_GATEWAY_URL"] = custom_url mock_session = MagicMock() mock_session.region_name = None # nothing configured anywhere - with patch("boto3.Session", return_value=mock_session): - with pytest.raises(ValueError, match="signing region"): - config.sign_request( - headers={}, - optional_params={}, - request_data={"jsonrpc": "2.0"}, - api_base="https://gateway.internal.example.com/mcp", - ) + try: + with patch("boto3.Session", return_value=mock_session): + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=custom_url, + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) def test_sign_request_custom_hostname_uses_shared_config_region(self): """Custom hostname + region from AWS shared config (profile) must be honored.""" config = AgentCoreSearchConfig() + custom_url = "https://gateway.internal.example.com/mcp" + os.environ["AGENTCORE_GATEWAY_URL"] = custom_url mock_session = MagicMock() mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile - with ( - patch("boto3.Session", return_value=mock_session), - patch.object( - AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM - "_sign_request", - return_value=({}, b"{}"), - ) as mock_base_sign, - ): - config.sign_request( - headers={}, - optional_params={}, - request_data={"jsonrpc": "2.0"}, - api_base="https://gateway.internal.example.com/mcp", - ) - assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + try: + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=custom_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" @@ -434,7 +523,11 @@ class TestAgentCoreSearchEdgeCases: assert getattr(err, "status_code", None) == 503 assert "boom" in str(err) - def test_search_cost_lookup_is_mapped(self): + def test_search_cost_lookup_is_mapped(self, monkeypatch): + """Assert against the map in this checkout: the remote cost map litellm loads by + default only carries providers already released.""" + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.search.cost_calculator import search_provider_cost_per_query + monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0) From 49dca4996c2096d8d5bba1ab5e1f10028f221abf Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 19 Aug 2026 19:31:19 +0000 Subject: [PATCH 27/88] fix(search): point the AgentCore region error at AWS_DEFAULT_REGION boto3's session resolution ignores AWS_REGION, so an operator following the old message still hit the same failure. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 85c5818031b..80e20d00ff9 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -331,7 +331,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): return configured_region raise ValueError( f"Cannot derive the SigV4 signing region from api_base '{api_base}' " - "or the AWS configuration chain. Set aws_region_name (or AWS_REGION / " + "or the AWS configuration chain. Set aws_region_name (or AWS_DEFAULT_REGION / " "a profile region) to the gateway's region when using a custom hostname." ) From a613773fca0161df9ab879f12f02caaa7464af99 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 19 Aug 2026 14:02:15 -0700 Subject: [PATCH 28/88] feat(auto-router)!: scope shadow eval jobs to multiple keys (#37251) * feat(auto-router): scope shadow eval jobs to multiple keys A shadow eval job now covers a set of keys instead of exactly one, and each key carries its own max_turns budget, so one key exhausting its budget leaves its siblings sampling. The existing job row already is the per-key unit (api_key_id, max_turns, stopped_at, and the one-active-per-key-and-direction partial unique index all live on it), so multi-key is grouping rather than schema surgery: a new group_id column ties N sibling rows written atomically by one create_many, the API's job id becomes the group id, and pre-existing jobs backfill group_id = id so their ids keep resolving. The sampler hot path is untouched; its test file has a zero-line diff Results come back pooled plus a per-key breakdown and responses list every key with its own budget, stop state and read-time labels. The dashboard is adapted minimally to the new shapes (the picker stays single-key and submits a one-key list); the multi-select picker and per-key table land in the stacked UI PR * fix(shadow_eval): derive completed from spent budgets and record operator stops * fix(shadow_eval): stamp stops atomically and freeze counts at the stamp The stop endpoint wrote stopped_by and stopped_at as two separate updates, so a failure between them left a job reading stopped while its unstamped legs kept sampling, and the retry got 400 already stopped. One UPDATE now stamps stopped_by and every missing stopped_at together, preserving the stopped_at a leg earned from its own budget via COALESCE Attempt counts now exclude attempts that land after a leg's stopped_at, so an in-flight attempt finishing just after an operator stop can never push a legacy pre-stopped_by job over its budget and flip it from stopped to completed at read time * fix(shadow_eval): backfill stopped_by so legacy stops never read as completions * chore(ui): regenerate api types for the shadow eval stop fields * fix(shadow_eval): let the stop statement pick one winner under racing stops Two operators can both pass the derived-status guard in the race window. The stop UPDATE now claims only legs with stopped_by still null and the endpoint judges by its row count, so exactly one caller ever gets the 200 and the loser gets the same already-stopped 400 a late caller gets * refactor(shadow_eval): make the stop statement the whole state machine The status guard ran before the UPDATE, so a stop racing the last budgeted attempt still claimed the job and it read stopped forever instead of completed. The statement now claims the job only while a leg still samples inside the window with no stop recorded, and the endpoint reads once after writing: a racing operator, a same-instant budget spend, and a repeat stop all get the 400 naming the status the job actually holds. The pre-write guard and the hand-built response go away * chore(ui): regenerate api types for the stop route description --- .../migration.sql | 7 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 30 +- .../auto_router_endpoints.py | 378 ++++++--- litellm/proxy/schema.prisma | 30 +- .../auto_router_endpoints.py | 103 ++- schema.prisma | 30 +- .../test_auto_router_endpoints.py | 717 ++++++++++++++---- .../_components/ShadowEvalSection.test.tsx | 43 +- .../_components/ShadowEvalSection.tsx | 26 +- .../_components/useShadowEval.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 118 ++- 12 files changed, 1117 insertions(+), 371 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260817000000_shadow_eval_multi_key/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260818224500_add_shadow_eval_stopped_by/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817000000_shadow_eval_multi_key/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817000000_shadow_eval_multi_key/migration.sql new file mode 100644 index 00000000000..18ef5c40662 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817000000_shadow_eval_multi_key/migration.sql @@ -0,0 +1,7 @@ +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "group_id" TEXT; + +UPDATE "LiteLLM_ShadowEvalJob" SET "group_id" = "id" WHERE "group_id" IS NULL; + +ALTER TABLE "LiteLLM_ShadowEvalJob" ALTER COLUMN "group_id" SET NOT NULL; + +CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_group_id_idx" ON "LiteLLM_ShadowEvalJob"("group_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818224500_add_shadow_eval_stopped_by/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818224500_add_shadow_eval_stopped_by/migration.sql new file mode 100644 index 00000000000..9efa3fdd052 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818224500_add_shadow_eval_stopped_by/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "stopped_by" TEXT; + +UPDATE "LiteLLM_ShadowEvalJob" SET stopped_by = 'unknown' +WHERE stopped_at IS NOT NULL AND ends_at > (NOW() AT TIME ZONE 'utc'); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 52fb447157b..f79e2bb0c18 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1467,28 +1467,38 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests -// the router did serve against a fixed baseline model, answering whether a key already on -// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in +// either direction. forward duplicates the requests the keys did not route through the +// router through it, answering whether they should adopt it; reverse duplicates the +// requests the router did serve against a fixed baseline model, answering whether a key +// already on it still benefits. Either way a sampled slice runs in a detached task and an +// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job: +// immutable config plus that key's own turn budget and stop state, so one key exhausting +// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id +// (the id the API reports), written together by one atomic create_many with identical +// config; single-key jobs predating group_id were backfilled group_id = id. "One active +// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE +// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state +// partial indexes; it is what makes a concurrent start on another pod race-safe rather +// than read-then-create. Every count, status, and spend figure is derived from the +// append-only attempt rows, so nothing can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed + group_id String // legs of one job share this; the API's job id + api_key_id String // hashed virtual key whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns + max_turns Int // this key's sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime stopped_at DateTime? + stopped_by String? // operator who stopped it early; null when it ended on its own + @@index([group_id]) @@index([api_key_id]) @@index([created_at]) } diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d8b7414f32c..0112ad1f6ed 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -6,10 +6,13 @@ POST /auto_router/test_routing - Route one prompt through an unsaved complexity- from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from itertools import groupby +from operator import attrgetter from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Protocol +from uuid import uuid4 -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, field_validator from litellm._logging import verbose_proxy_logger from litellm.exceptions import BudgetExceededError @@ -41,6 +44,8 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterRoutingTestRequest, AutoRouterRoutingTestResponse, RequestComplexityRouterConfig, + ShadowEvalDirection, + ShadowEvalJobKeyResponse, ShadowEvalJobResponse, ShadowEvalResult, ShadowEvalSlice, @@ -89,17 +94,9 @@ class _ShadowEvalJobRow(Protocol): class _ShadowEvalJobTable(Protocol): - async def find_unique(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ShadowEvalJobRow]: ... - async def find_first(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ... - - async def find_many( - self, *, where: Mapping[str, object], order: Mapping[str, str], take: int - ) -> Sequence[_ShadowEvalJobRow]: ... - - async def create(self, data: Mapping[str, object]) -> _ShadowEvalJobRow: ... - - async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _ShadowEvalJobRow | None: ... + async def create_many(self, data: Sequence[Mapping[str, object]]) -> int: ... class _ShadowEvalAttemptRow(Protocol): @@ -606,18 +603,19 @@ _ATTEMPT_AGG_SELECT: Final = """ COUNT(*) FILTER (WHERE outcome = 'tie')::int AS ties, AVG(confidence)::float AS avg_confidence FROM "LiteLLM_ShadowEvalAttempt" -WHERE job_id = $1 AND outcome != 'error' +WHERE job_id = ANY($1::text[]) AND outcome != 'error' GROUP BY 1 """ _ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT +_ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT _SWEEP_FINISHED_JOBS_SQL: Final = """ -UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = NOW() -WHERE j.api_key_id = $1 AND j.stopped_at IS NULL +UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = (NOW() AT TIME ZONE 'utc') +WHERE j.api_key_id = ANY($1::text[]) AND j.stopped_at IS NULL AND ( - j.ends_at <= NOW() + j.ends_at <= (NOW() AT TIME ZONE 'utc') OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns ) """ @@ -628,7 +626,52 @@ SELECT COUNT(*) FILTER (WHERE outcome = 'error')::int AS error_count, COALESCE(SUM(judge_cost), 0)::float AS judge_spend FROM "LiteLLM_ShadowEvalAttempt" -WHERE job_id = $1 +WHERE job_id = ANY($1::text[]) +""" + +_ATTEMPT_COUNTS_SQL: Final = """ +SELECT a.job_id, COUNT(*)::int AS attempt_count +FROM "LiteLLM_ShadowEvalAttempt" a +JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id +WHERE a.job_id = ANY($1::text[]) AND (j.stopped_at IS NULL OR a.created_at <= j.stopped_at) +GROUP BY a.job_id +""" + +_STOP_JOB_SQL: Final = """ +UPDATE "LiteLLM_ShadowEvalJob" +SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp) +WHERE group_id = $1 AND stopped_by IS NULL + AND ends_at > (NOW() AT TIME ZONE 'utc') + AND EXISTS ( + SELECT 1 FROM "LiteLLM_ShadowEvalJob" k + WHERE k.group_id = $1 AND k.stopped_at IS NULL + AND (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = k.id) < k.max_turns + ) +""" + + +class _AttemptCountRow(BaseModel): + job_id: str + attempt_count: int + + +_ATTEMPT_COUNT_ROWS: Final = TypeAdapter(list[_AttemptCountRow]) + + +_LIST_LEGS_SQL: Final = """ +SELECT * FROM "LiteLLM_ShadowEvalJob" +WHERE group_id IN ( + SELECT group_id FROM "LiteLLM_ShadowEvalJob" + GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int +) +""" + +_LIST_LEGS_BY_KEY_SQL: Final = """ +SELECT * FROM "LiteLLM_ShadowEvalJob" +WHERE group_id IN ( + SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE api_key_id = $2 + GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int +) """ @@ -659,18 +702,98 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]: ) +class _LegRow(BaseModel): + """One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is + one key's leg of a job; the legs of a job share group_id and identical config, written + together by one create_many. The API's job id is the group id, so leg ids never leave + the server (attempts reference them internally).""" + + model_config = ConfigDict(from_attributes=True) + + id: str + group_id: str + api_key_id: str + router_name: str + direction: ShadowEvalDirection + baseline_model: str | None = None + judge_model: str + shadow_percentage: float + max_turns: int + created_at: datetime + ends_at: datetime + stopped_at: datetime | None = None + stopped_by: str | None = None + + @field_validator("created_at", "ends_at", "stopped_at") + @classmethod + def _as_aware_utc(cls, value: datetime | None) -> datetime | None: + """The columns store naive UTC wall time (prisma's convention); prisma reads hand + back aware datetimes while raw SQL reads hand back naive ones, so this boundary + makes every read aware UTC before anything compares or serializes them.""" + if value is None or value.tzinfo is not None: + return value + return value.replace(tzinfo=timezone.utc) + + +_LEG_ROWS: Final = TypeAdapter(list[_LegRow]) + + +async def _leg_attempt_counts(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> Mapping[str, int]: + """Each leg's attempt count by leg id, judged and errored alike, in one grouped read. + It is the same count the sampler budgets against max_turns, so the derived status + flips to completed exactly when sampling actually ends. A stamped leg's count freezes + at its stopped_at: in-flight attempts that land after the stamp are excluded, so they + can never reclassify a leg that was stopped under budget as budget-spent.""" + if not legs: + return MappingProxyType({}) + rows: Final = _ATTEMPT_COUNT_ROWS.validate_python( + await _query_raw(prisma_client, _ATTEMPT_COUNTS_SQL, [leg.id for leg in legs]) # mutable-ok: query param + or () + ) + return MappingProxyType({row.job_id: row.attempt_count for row in rows}) + + +def _group_response(group_id: str, legs: Sequence[_LegRow], attempt_counts: Mapping[str, int]) -> ShadowEvalJobResponse: + """The one constructor of a job response: the caller names the group and passes that + group's legs. Config is read off the first leg because every leg carries the same copy, + written by one create_many. No caller may serialize a raw row (that would leak a leg id + as the job id).""" + first: Final = legs[0] + return ShadowEvalJobResponse( + job_id=group_id, + keys=tuple( + ShadowEvalJobKeyResponse( + api_key_id=leg.api_key_id, + max_turns=leg.max_turns, + stopped_at=leg.stopped_at, + attempt_count=attempt_counts.get(leg.id, 0), + ) + for leg in sorted(legs, key=lambda leg: leg.api_key_id) + ), + router_name=first.router_name, + direction=first.direction, + baseline_model=first.baseline_model, + judge_model=first.judge_model, + shadow_percentage=first.shadow_percentage, + created_at=first.created_at, + ends_at=first.ends_at, + stopped_by=next((leg.stopped_by for leg in legs if leg.stopped_by is not None), None), + ) + + _NO_KEY_LABELS: Final[tuple[str | None, str | None]] = (None, None) async def _with_key_labels( prisma_client: "PrismaClient", responses: Sequence[ShadowEvalJobResponse] ) -> tuple[ShadowEvalJobResponse, ...]: - """Resolve each job's key hash to the key's alias and masked name in one batched read, + """Resolve every scoped key's hash to its alias and masked name in one batched read, so the UI can say whose traffic a job shadows. Deleted keys resolve to None.""" if not responses: return () + tokens: Final = sorted(frozenset(key.api_key_id for response in responses for key in response.keys)) key_rows: Final = await _verification_tokens(prisma_client).find_many( - where={"token": {"in": sorted({response.api_key_id for response in responses})}} # mutable-ok: Prisma filter + where={"token": {"in": tokens}} # mutable-ok: Prisma filter ) labels: Final[Mapping[str, tuple[str | None, str | None]]] = { row.token: (row.key_alias, row.key_name) for row in key_rows or () @@ -678,32 +801,50 @@ async def _with_key_labels( return tuple( response.model_copy( update={ # mutable-ok: pydantic update payload - "key_alias": labels.get(response.api_key_id, _NO_KEY_LABELS)[0], - "key_name": labels.get(response.api_key_id, _NO_KEY_LABELS)[1], + "keys": tuple( + key.model_copy( + update={ # mutable-ok: pydantic update payload + "key_alias": labels.get(key.api_key_id, _NO_KEY_LABELS)[0], + "key_name": labels.get(key.api_key_id, _NO_KEY_LABELS)[1], + } + ) + for key in response.keys + ) } ) for response in responses ) -async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None: - """Both stratifications of one job's verdicts. Tier answers "where does the router do - well"; the model stratification groups by whichever model served the real arm, so it - answers "which of the models this key uses today would the router beat" forward, and - "for the turns the router sent to X, did X beat the baseline" in reverse. Reads are - bounded by the job's own attempts (<= max_turns) via the job_id index.""" +async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> ShadowEvalResult | None: + """All three stratifications of one job's verdicts. Tier answers "where does the router + do well"; the model stratification groups by whichever model served the real arm, so it + answers "which of the models these keys use today would the router beat" forward, and + "for the turns the router sent to X, did X beat the baseline" in reverse; key answers + "which key's traffic does the router suit". Reads are bounded by the job's own attempts + (<= the sum of its keys' max_turns) via the job_id index.""" + leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python( - await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, job_id) or () + await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, leg_ids) or () ) if not by_tier: return None by_model: Final = _ATTEMPT_AGG_ROWS.validate_python( - await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, job_id) or () + await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, leg_ids) or () + ) + key_by_leg: Final = MappingProxyType({leg.id: leg.api_key_id for leg in legs}) + by_leg: Final = _ATTEMPT_AGG_ROWS.validate_python( + await _query_raw(prisma_client, _ATTEMPT_AGG_BY_LEG_SQL, leg_ids) or () + ) + by_key: Final = tuple( + row.model_copy(update={"grp": key_by_leg[row.grp]}) # mutable-ok: pydantic update payload + for row in by_leg ) total_turns: Final = sum(r.turn_count for r in by_tier) return ShadowEvalResult( by_tier=_slices(by_tier), by_current_model=_slices(by_model), + by_key=_slices(by_key), overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns), overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns), ) @@ -721,20 +862,21 @@ async def start_shadow_eval( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: """ - Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second - arm, judge the two responses blind, and stratify win rates by tier and by the model that - served the real arm. + Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against + a second arm, judge the two responses blind, and stratify win rates by tier, by the model + that served the real arm, and by key. - A forward job answers whether the key should adopt router_name: it samples the requests + A forward job answers whether the keys should adopt router_name: it samples the requests the router did not serve and duplicates them through it. A reverse job answers whether a key already on the router still gains from it: it samples the requests the router did serve and duplicates them against baseline_model. A key can hold one active job per direction, so both questions can run at once. - Shadow responses are never served to users. The job samples until it has judged - max_turns turns, reaches the end of its window, or is stopped; sampling changes - propagate to pods within about 10 seconds. Shadow and judge calls bill to the - shadowed key but are excluded from request counts and auto-router adoption metrics. + Shadow responses are never served to users. Each key samples until it has judged + max_turns turns of its own traffic, the job's window ends, or the job is stopped, so one + key running out of budget does not end sampling for the others; sampling changes + propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed + key but are excluded from request counts and auto-router adoption metrics. """ from litellm.proxy.proxy_server import llm_router, prisma_client @@ -746,48 +888,58 @@ async def start_shadow_eval( _validate_plain_model(llm_router, data.judge_model, "judge_model") if data.baseline_model is not None: _validate_plain_model(llm_router, data.baseline_model, "baseline_model") - key_row: Final = await _verification_tokens(prisma_client).find_unique( - where={"token": data.api_key_id} # mutable-ok: Prisma filter + token_rows: Final = await _verification_tokens(prisma_client).find_many( + where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter ) - if key_row is None: + unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ()))) + if unknown: raise HTTPException( status_code=400, detail=( - f"api_key_id '{data.api_key_id}' is not a key on this proxy; pass the key's token hash, " + f"api_key_ids not on this proxy: {', '.join(unknown)}; pass each key's token hash, " "the value the key list and key info endpoints report" ), ) - # A job that expired or exhausted its turn budget stopped sampling on its own, but - # still holds its slot in the per-key, per-direction partial unique index until - # stamped; free it so a new eval can start. Sweeping both directions is deliberate. - await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id) - active: Final = await _shadow_eval_jobs(prisma_client).find_first( + # A job whose window passed or whose turn budget ran out stopped sampling on its own, + # but its legs still hold their slots in the per-key, per-direction partial unique index + # until stamped; free them so a new eval can start. Sweeping both directions is deliberate. + requested: Final = list(data.api_key_ids) # mutable-ok: query param + await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, requested) + claimed: Final = await _shadow_eval_jobs(prisma_client).find_many( where={ # mutable-ok: Prisma filter - "api_key_id": data.api_key_id, + "api_key_id": {"in": requested}, # mutable-ok: Prisma filter "direction": data.direction, "stopped_at": None, }, ) - if active is not None: + if claimed: raise HTTPException( status_code=409, - detail=f"Key already has an active {data.direction} shadow eval job ({active.id}). Stop it first.", + detail=( + f"Already in an active {data.direction} shadow eval job: " + + ", ".join(sorted(f"{row.api_key_id} (job {row.group_id})" for row in claimed)) + + ". Stop it first." + ), ) now: Final = datetime.now(timezone.utc) + group_id: Final = str(uuid4()) + ends_at: Final = now + timedelta(days=data.duration_days) + shared_config: Final = { # mutable-ok: Prisma payload + "group_id": group_id, + "router_name": data.router_name, + "direction": data.direction, + "baseline_model": data.baseline_model, + "judge_model": data.judge_model, + "shadow_percentage": data.shadow_percentage, + "max_turns": data.max_turns, + "created_by": user_api_key_dict.user_id, + "created_at": now, + "ends_at": ends_at, + } try: - job: Final = await _shadow_eval_jobs(prisma_client).create( - data={ # mutable-ok: Prisma payload - "api_key_id": data.api_key_id, - "router_name": data.router_name, - "direction": data.direction, - "baseline_model": data.baseline_model, - "judge_model": data.judge_model, - "shadow_percentage": data.shadow_percentage, - "max_turns": data.max_turns, - "created_by": user_api_key_dict.user_id, - "ends_at": now + timedelta(days=data.duration_days), - } + await _shadow_eval_jobs(prisma_client).create_many( + data=[{**shared_config, "api_key_id": key} for key in data.api_key_ids] # mutable-ok: Prisma payload ) except Exception as e: if not _is_unique_violation(e): @@ -795,11 +947,28 @@ async def start_shadow_eval( raise HTTPException( status_code=409, detail=( - f"Key already has an active {data.direction} shadow eval job (started concurrently). Stop it first." + f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first." ), ) from e - return ShadowEvalJobResponse.model_validate(job, from_attributes=True).model_copy( - update={"key_alias": key_row.key_alias, "key_name": key_row.key_name} # mutable-ok: pydantic update payload + labels: Final = MappingProxyType({row.token: row for row in token_rows}) + return ShadowEvalJobResponse( + job_id=group_id, + keys=tuple( + ShadowEvalJobKeyResponse( + api_key_id=api_key_id, + max_turns=data.max_turns, + key_alias=labels[api_key_id].key_alias, + key_name=labels[api_key_id].key_name, + ) + for api_key_id in sorted(data.api_key_ids) + ), + router_name=data.router_name, + direction=data.direction, + baseline_model=data.baseline_model, + judge_model=data.judge_model, + shadow_percentage=data.shadow_percentage, + created_at=now, + ends_at=ends_at, ) @@ -811,23 +980,38 @@ async def start_shadow_eval( ) async def list_shadow_eval_jobs( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], - api_key_id: Annotated[str | None, Query(description="Filter to jobs shadowing this key")] = None, + api_key_id: Annotated[ + str | None, Query(description="Filter to jobs that shadow this key, alone or alongside others") + ] = None, limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50, ) -> tuple[ShadowEvalJobResponse, ...]: - """List shadow eval jobs, newest first. Counts and results ride the detail endpoint only.""" + """List shadow eval jobs, newest first, each key with its attempt count so status is + accurate. Judged counts, spend, and results ride the detail endpoint only.""" from litellm.proxy.proxy_server import prisma_client _require_admin_viewer(user_api_key_dict, "view shadow evals") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - records: Final = await _shadow_eval_jobs(prisma_client).find_many( - where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter - order={"created_at": "desc"}, # mutable-ok: Prisma order - take=limit, + legs: Final = _LEG_ROWS.validate_python( + ( + await _query_raw(prisma_client, _LIST_LEGS_BY_KEY_SQL, limit, api_key_id) + if api_key_id + else await _query_raw(prisma_client, _LIST_LEGS_SQL, limit) + ) + or () ) + by_group: Final[Mapping[str, tuple[_LegRow, ...]]] = MappingProxyType( + { + group_id: tuple(group) + for group_id, group in groupby(sorted(legs, key=attrgetter("group_id")), key=attrgetter("group_id")) + } + ) + newest_first: Final = sorted( + by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True + ) + counts: Final = await _leg_attempt_counts(prisma_client, legs) return await _with_key_labels( - prisma_client, - tuple(ShadowEvalJobResponse.model_validate(record, from_attributes=True) for record in records or ()), + prisma_client, tuple(_group_response(group_id, by_group[group_id], counts) for group_id in newest_first) ) @@ -847,20 +1031,24 @@ async def get_shadow_eval_job( _require_admin_viewer(user_api_key_dict, "view shadow evals") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - record: Final = await _shadow_eval_jobs(prisma_client).find_unique( - where={"id": job_id} # mutable-ok: Prisma filter + legs: Final = _LEG_ROWS.validate_python( + await _shadow_eval_jobs(prisma_client).find_many( + where={"group_id": job_id} # mutable-ok: Prisma filter + ) + or () ) - if record is None: + if not legs: raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}") + leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param totals: Final = _ATTEMPT_TOTALS_ROWS.validate_python( - await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, job_id) or () + await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, leg_ids) or () ) latest_error: Final = await _shadow_eval_attempts(prisma_client).find_first( - where={"job_id": job_id, "outcome": "error"}, # mutable-ok: Prisma filter + where={"job_id": {"in": leg_ids}, "outcome": "error"}, # mutable-ok: Prisma filter order={"created_at": "desc"}, # mutable-ok: Prisma order ) labeled: Final = await _with_key_labels( - prisma_client, (ShadowEvalJobResponse.model_validate(record, from_attributes=True),) + prisma_client, (_group_response(job_id, legs, await _leg_attempt_counts(prisma_client, legs)),) ) return labeled[0].model_copy( update={ # mutable-ok: pydantic update payload @@ -868,7 +1056,7 @@ async def get_shadow_eval_job( "error_count": totals[0].error_count if totals else 0, "judge_spend": round(totals[0].judge_spend, 6) if totals else 0.0, "last_error": latest_error.error if latest_error else None, - "results": await _shadow_eval_results(prisma_client, job_id), + "results": await _shadow_eval_results(prisma_client, legs), } ) @@ -883,25 +1071,33 @@ async def stop_shadow_eval_job( job_id: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: - """Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s.""" + """Stop an active shadow eval job, every key it scopes at once. Attempts are kept; + sampling halts within ~10s. Keys that already stopped on their own budget keep the + stopped_at they earned. The statement is the whole state machine: it claims the job + only while a leg still samples inside the window with no stop recorded, so a racing + operator, a same-instant budget spend, and a repeat stop all read the same 400 with + the status the job actually holds.""" from litellm.proxy.proxy_server import prisma_client _require_admin_writer(user_api_key_dict, "stop a shadow eval") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - record: Final = await _shadow_eval_jobs(prisma_client).find_unique( - where={"id": job_id} # mutable-ok: Prisma filter + stamp: Final = datetime.now(timezone.utc) + operator: Final = user_api_key_dict.user_id or "operator" + claimed: Final = await prisma_client.db.execute_raw( + _STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat() ) - if record is None: + legs: Final = _LEG_ROWS.validate_python( + await _shadow_eval_jobs(prisma_client).find_many( + where={"group_id": job_id} # mutable-ok: Prisma filter + ) + or () + ) + if not legs: raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}") - current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True) - if current.status != "running": + counts: Final = await _leg_attempt_counts(prisma_client, legs) + current: Final = _group_response(job_id, legs, counts) + if claimed == 0: raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}") - updated: Final = await _shadow_eval_jobs(prisma_client).update( - where={"id": job_id}, # mutable-ok: Prisma filter - data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload - ) - labeled: Final = await _with_key_labels( - prisma_client, (ShadowEvalJobResponse.model_validate(updated, from_attributes=True),) - ) + labeled: Final = await _with_key_labels(prisma_client, (current,)) return labeled[0] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 52fb447157b..f79e2bb0c18 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1467,28 +1467,38 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests -// the router did serve against a fixed baseline model, answering whether a key already on -// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in +// either direction. forward duplicates the requests the keys did not route through the +// router through it, answering whether they should adopt it; reverse duplicates the +// requests the router did serve against a fixed baseline model, answering whether a key +// already on it still benefits. Either way a sampled slice runs in a detached task and an +// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job: +// immutable config plus that key's own turn budget and stop state, so one key exhausting +// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id +// (the id the API reports), written together by one atomic create_many with identical +// config; single-key jobs predating group_id were backfilled group_id = id. "One active +// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE +// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state +// partial indexes; it is what makes a concurrent start on another pod race-safe rather +// than read-then-create. Every count, status, and spend figure is derived from the +// append-only attempt rows, so nothing can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed + group_id String // legs of one job share this; the API's job id + api_key_id String // hashed virtual key whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns + max_turns Int // this key's sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime stopped_at DateTime? + stopped_by String? // operator who stopped it early; null when it ended on its own + @@index([group_id]) @@index([api_key_id]) @@index([created_at]) } diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 9461297feca..63c93e0f268 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from datetime import datetime, timezone from typing import Final, Literal, TypeAlias -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, computed_field, field_validator, model_validator +from pydantic import BaseModel, Field, computed_field, field_validator, model_validator from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig from litellm.types.utils import StandardLoggingRoutingDecision @@ -155,13 +155,17 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" class StartShadowEvalRequest(BaseModel): - """Start duplicating a key's traffic for blind comparison against an auto-router.""" + """Start duplicating one or more keys' traffic for blind comparison against an auto-router.""" - api_key_id: str = Field( + api_key_ids: tuple[str, ...] = Field( + min_length=1, + max_length=100, description=( - "The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this " - "key's traffic; requests made with any other key are not sampled." - ) + "The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these " + "keys' traffic; requests made with any other key are not sampled. Each key carries its own " + "max_turns budget, so one key exhausting its budget leaves the others sampling. At most 100 " + "keys per job, which also bounds every read the job's endpoints make." + ), ) router_name: str = Field(description="The auto-router under evaluation, in either direction") direction: ShadowEvalDirection = Field( @@ -204,8 +208,9 @@ class StartShadowEvalRequest(BaseModel): ge=1, le=2000, description=( - "Sample budget: the job judges at most this many turns, then completes. This is also the spend " - "bound; expected judge cost is roughly max_turns times one judge call" + "Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, " + "so a job over N keys judges at most N times max_turns turns. This is also the spend bound; " + "expected judge cost is roughly that turn ceiling times one judge call" ), ) @@ -214,6 +219,12 @@ class StartShadowEvalRequest(BaseModel): def _round_percentage(cls, value: float) -> float: return round(value, 2) + @field_validator("api_key_ids") + @classmethod + def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """A key named twice would collide with itself on the one-active-per-(key, direction) index.""" + return tuple(dict.fromkeys(value)) + @model_validator(mode="after") def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest": if self.direction == "reverse" and self.baseline_model is None: @@ -251,24 +262,46 @@ class ShadowEvalResult(BaseModel): by_tier: tuple[ShadowEvalSlice, ...] by_current_model: tuple[ShadowEvalSlice, ...] = Field( description=( - "Sliced by the model that served the real arm: the key's incumbent models in forward mode, " + "Sliced by the model that served the real arm: the keys' incumbent models in forward mode, " "and in reverse the models the router itself picked" ) ) + by_key: tuple[ShadowEvalSlice, ...] = Field( + description=( + "One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job " + "scopes but has not judged a turn for yet are absent rather than reported as zero" + ), + ) overall_shadow_win_rate_pct: float overall_tie_rate_pct: float -class ShadowEvalJobResponse(BaseModel): - """A shadow-eval job. Validates directly from the prisma record (job_id reads the - row's id); status is derived from stopped_at and ends_at, never stored, so no writer - anywhere can produce an inconsistent one. Aggregate fields are populated by the - detail endpoint only and stay None on list responses.""" +class ShadowEvalJobKeyResponse(BaseModel): + """One key a job shadows, with its own budget and stop state.""" - model_config = ConfigDict(from_attributes=True, populate_by_name=True) + api_key_id: str = Field(description="The hashed virtual key whose traffic this entry scopes") + max_turns: int = Field(description="This key's own sample budget, independent of its siblings'") + stopped_at: datetime | None = Field( + default=None, + description=( + "When this key's slot was stamped free, whether its own budget ran out, the window closed, " + "or an operator stopped the job; status is derived, so a spent budget reads completed even " + "while this is still unset" + ), + ) + attempt_count: int | None = Field( + default=None, + description=( + "This key's sampled attempts so far, judged and errored alike, the same count the sampler " + "budgets against max_turns; populated on list and detail responses. Frozen at stopped_at " + "once the key is stamped, so in-flight attempts landing after a stop never reclassify it" + ), + ) + + @property + def budget_spent(self) -> bool: + return self.attempt_count is not None and self.attempt_count >= self.max_turns - job_id: str = Field(validation_alias=AliasChoices("id", "job_id")) - api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's") key_alias: str | None = Field( default=None, description="Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted", @@ -277,15 +310,34 @@ class ShadowEvalJobResponse(BaseModel): default=None, description="Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias", ) + + +class ShadowEvalJobResponse(BaseModel): + """A shadow-eval job over one or more keys, each with its own budget and stop state; + status is derived from stopped_by, the keys' stop and budget state, and ends_at, + never stored, so no writer anywhere can produce an inconsistent one. Aggregate + fields are populated by the detail endpoint only and stay None on list responses.""" + + job_id: str + keys: tuple[ShadowEvalJobKeyResponse, ...] = Field( + min_length=1, + description="The keys whose traffic this job evaluates, and only those keys', each with its own budget", + ) router_name: str direction: ShadowEvalDirection = "forward" baseline_model: str | None = None judge_model: str shadow_percentage: float - max_turns: int created_at: datetime ends_at: datetime - stopped_at: datetime | None = None + stopped_by: str | None = Field( + default=None, + description=( + "The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled " + "by migration for jobs that displayed stopped when the column arrived; None when the job " + "ended on its own. Its presence is what makes a job read stopped rather than completed" + ), + ) judged_count: int | None = Field(default=None, description="Verdicts recorded; detail endpoint only") error_count: int | None = Field(default=None, description="Sampled attempts that errored; detail endpoint only") @@ -296,12 +348,19 @@ class ShadowEvalJobResponse(BaseModel): @computed_field @property def status(self) -> ShadowEvalStatus: - """A job whose window has passed reads completed even if a later sweep stamped - stopped_at; stopped means sampling ended before the window did.""" + """Three recorded facts, no history-guessing: a stop is stopped_by (the migration + backfills it for every job that displayed stopped when the column arrived, so the + pre-column population is closed), completion is the window passing or every key + spending its budget, and anything else is running. The all-keys-stamped fallback + covers only stops written by pre-column pods during a rolling deploy.""" + if self.stopped_by is not None: + return "stopped" if datetime.now(timezone.utc) >= ( self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc) ): return "completed" - if self.stopped_at is not None: + if all(key.budget_spent for key in self.keys): + return "completed" + if all(key.stopped_at is not None for key in self.keys): return "stopped" return "running" diff --git a/schema.prisma b/schema.prisma index 52fb447157b..f79e2bb0c18 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1467,28 +1467,38 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests -// the router did serve against a fixed baseline model, answering whether a key already on -// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in +// either direction. forward duplicates the requests the keys did not route through the +// router through it, answering whether they should adopt it; reverse duplicates the +// requests the router did serve against a fixed baseline model, answering whether a key +// already on it still benefits. Either way a sampled slice runs in a detached task and an +// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job: +// immutable config plus that key's own turn budget and stop state, so one key exhausting +// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id +// (the id the API reports), written together by one atomic create_many with identical +// config; single-key jobs predating group_id were backfilled group_id = id. "One active +// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE +// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state +// partial indexes; it is what makes a concurrent start on another pod race-safe rather +// than read-then-create. Every count, status, and spend figure is derived from the +// append-only attempt rows, so nothing can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed + group_id String // legs of one job share this; the API's job id + api_key_id String // hashed virtual key whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns + max_turns Int // this key's sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime stopped_at DateTime? + stopped_by String? // operator who stopped it early; null when it ended on its own + @@index([group_id]) @@index([api_key_id]) @@index([created_at]) } diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 77149457e82..3fd023552f5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -4,6 +4,7 @@ Unit tests for auto router management endpoints import os import sys +from pathlib import Path import pytest from fastapi import HTTPException @@ -325,9 +326,7 @@ class TestAutoRouterBenchmarks: from litellm.proxy.management_endpoints.auto_router_endpoints import _benchmark_totals totals = _benchmark_totals(self.ROW) - bucket_hits = ( - totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits - ) + bucket_hits = totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits assert bucket_hits == 27 assert totals.cache.hit_rate_pct == pytest.approx(100.0 * 28 / 38, abs=0.1) @@ -490,7 +489,7 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( start_shadow_eval, stop_shadow_eval_job, ) -from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalJobResponse, StartShadowEvalRequest +from litellm.types.management_endpoints.auto_router_endpoints import StartShadowEvalRequest VIEWER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, api_key="sk-view", user_id="viewer") NON_ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="user") @@ -507,19 +506,23 @@ def _shadow_router() -> MagicMock: return router -def _job_record(**overrides: object) -> MagicMock: - """Spec'd like a real prisma row: only the table's columns exist as attributes, so - from_attributes validation falls back to model defaults for everything else.""" +def _leg_record(**overrides: object) -> MagicMock: + """Spec'd like a real prisma row: only the table's columns exist as attributes. One + row is one key's leg of a job; legs sharing group_id are one job.""" defaults = { - "id": "job-1", + "id": "leg-1", + "group_id": "job-1", "api_key_id": "key-hash", "router_name": "my-router", + "direction": "forward", + "baseline_model": None, "judge_model": "anthropic/claude-sonnet-5", "shadow_percentage": 10.0, "max_turns": 200, "created_at": datetime(2026, 8, 11, tzinfo=timezone.utc), "ends_at": datetime.now(timezone.utc) + timedelta(days=7), "stopped_at": None, + "stopped_by": None, } fields = {**defaults, **overrides} record = MagicMock(spec=list(fields)) @@ -538,23 +541,90 @@ def _key_record( return record -def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock: +def _shadow_prisma(legs=(), agg_rows=None, by_leg_rows=None, known_keys=("key-hash", "key-hash-2")) -> MagicMock: + """The job-table fake honours the filters it is handed, so a read that forgets + stopped_at sees rows the partial index would have released, one that forgets + direction sees the opposite-direction legs a key may hold at the same time, and a + group read that matched on a leg id would come back empty.""" prisma = MagicMock() - prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=_key_record()) - prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[_key_record()]) - prisma.db.execute_raw = AsyncMock(return_value=0) - prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=active_job) - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=None) - prisma.db.litellm_shadowevaljob.find_many = AsyncMock(return_value=[]) - prisma.db.litellm_shadowevaljob.create = AsyncMock(return_value=_job_record()) - prisma.db.litellm_shadowevaljob.update = AsyncMock( - return_value=_job_record(stopped_at=datetime.now(timezone.utc)) - ) + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[_key_record(token) for token in known_keys]) + async def execute_raw(sql: str, *params: object): + if "SET stopped_by" in sql: + group = [row for row in stored if row.group_id == params[0]] + counts = {row["job_id"]: row["attempt_count"] for row in prisma.attempt_rows} + sampling = any(row.stopped_at is None and counts.get(row.id, 0) < row.max_turns for row in group) + window_open = bool(group) and group[0].ends_at > datetime.now(timezone.utc) + claimable = [row for row in group if row.stopped_by is None] + if not (claimable and sampling and window_open): + return 0 + for row in claimable: + row.stopped_by = params[1] + if row.stopped_at is None: + row.stopped_at = datetime.fromisoformat(str(params[2])).replace(tzinfo=timezone.utc) + return len(claimable) + return 0 + + prisma.db.execute_raw = AsyncMock(side_effect=execute_raw) + stored = legs if isinstance(legs, list) else list(legs) + + async def find_many_legs(where=None, **_: object): + current = list(stored) + w = dict(where or {}) + if "api_key_id" in w: + wanted = w["api_key_id"]["in"] if isinstance(w["api_key_id"], dict) else [w["api_key_id"]] + current = [row for row in current if row.api_key_id in wanted] + if "direction" in w: + current = [row for row in current if row.direction == w["direction"]] + if "stopped_at" in w: + current = [row for row in current if row.stopped_at is w["stopped_at"]] + if "group_id" in w: + wanted = w["group_id"]["in"] if isinstance(w["group_id"], dict) else [w["group_id"]] + current = [row for row in current if row.group_id in wanted] + return current + + def newest_groups(rows, limit): + latest: dict = {} + for row in rows: + if row.group_id not in latest or row.created_at > latest[row.group_id]: + latest[row.group_id] = row.created_at + ordered = sorted(latest, key=lambda group_id: latest[group_id], reverse=True) + return ordered[: int(limit)] + + def leg_dict(row): + fields = ( + "id", + "group_id", + "api_key_id", + "router_name", + "direction", + "baseline_model", + "judge_model", + "shadow_percentage", + "max_turns", + "created_at", + "ends_at", + "stopped_at", + "stopped_by", + ) + return {field: getattr(row, field) for field in fields} + + prisma.db.litellm_shadowevaljob.find_many = AsyncMock(side_effect=find_many_legs) + prisma.db.litellm_shadowevaljob.create_many = AsyncMock(return_value=1) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock(return_value=1) prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=None) + prisma.attempt_rows = [] async def query_raw(sql: str, *params: object): + if "AS attempt_count" in sql: + return prisma.attempt_rows + if "GROUP BY group_id" in sql: + scoped = [row for row in stored if "api_key_id = $2" not in sql or row.api_key_id == params[1]] + keep = set(newest_groups(scoped, params[0])) + return [leg_dict(row) for row in stored if row.group_id in keep] if "FILTER (WHERE outcome != 'error')::int AS judged_count" in sql: return [{"judged_count": 10, "error_count": 2, "judge_spend": 0.031}] + if "SELECT job_id AS grp" in sql: + return by_leg_rows if by_leg_rows is not None else [] return agg_rows if agg_rows is not None else [] prisma.db.query_raw = AsyncMock(side_effect=query_raw) @@ -563,7 +633,7 @@ def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock: def _start_request(**overrides: object) -> StartShadowEvalRequest: payload = { - "api_key_id": "key-hash", + "api_key_ids": ("key-hash",), "router_name": "my-router", "shadow_percentage": 10.0, "judge_model": "anthropic/claude-sonnet-5", @@ -575,44 +645,55 @@ def _start_request(**overrides: object) -> StartShadowEvalRequest: @pytest.mark.asyncio -async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones(monkeypatch: pytest.MonkeyPatch): - """Expiry and turn-budget exhaustion both end sampling on their own; either must - release the key's slot in the active-job index so a new eval can start.""" +async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeypatch: pytest.MonkeyPatch): + """N keys become N sibling rows sharing group_id and identical config, written by a + single create_many so a unique-index loser rolls back the whole claim, and expiry or + budget exhaustion frees every requested key's slot first.""" import litellm.proxy.proxy_server as proxy_server prisma = _shadow_prisma() monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - response = await start_shadow_eval(_start_request(), ADMIN) + response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) - assert response.status == "running" - assert response.max_turns == 200 - assert response.judged_count is None - sweep_sql, sweep_key = prisma.db.execute_raw.call_args.args + sweep_sql, sweep_keys = prisma.db.execute_raw.call_args.args assert "stopped_at IS NULL" in sweep_sql - assert "ends_at <= NOW()" in sweep_sql + assert "j.ends_at <= (NOW() AT TIME ZONE 'utc')" in sweep_sql + assert "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in sweep_sql assert ">= j.max_turns" in sweep_sql - assert sweep_key == "key-hash" - create_data = prisma.db.litellm_shadowevaljob.create.call_args.kwargs["data"] - assert create_data["api_key_id"] == "key-hash" - assert create_data["created_by"] == "admin" - assert "status" not in create_data + assert "j.api_key_id = ANY($1::text[])" in sweep_sql + assert sweep_keys == ["key-hash", "key-hash-2"] + prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert [row["api_key_id"] for row in rows] == ["key-hash", "key-hash-2"] + assert len({frozenset((k, v) for k, v in row.items() if k != "api_key_id") for row in rows}) == 1 + assert len({row["group_id"] for row in rows}) == 1 + assert all(row["max_turns"] == 200 and row["created_by"] == "admin" for row in rows) + assert all("status" not in row and "id" not in row for row in rows) + assert response.job_id == rows[0]["group_id"] + assert response.status == "running" + assert response.judged_count is None + assert [(key.api_key_id, key.max_turns, key.key_alias) for key in response.keys] == [ + ("key-hash", 200, "prod-alpha"), + ("key-hash-2", 200, "prod-alpha"), + ] @pytest.mark.asyncio @pytest.mark.parametrize( - "caller,request_overrides,active,expected_status", + "caller,request_overrides,claimed,expected_status", [ - (NON_ADMIN, {}, None, 403), - (VIEWER, {}, None, 403), - (ADMIN, {"router_name": "not-a-router"}, None, 400), - (ADMIN, {"judge_model": "not/a real model!"}, None, 400), - (ADMIN, {"judge_model": "my-router"}, None, 400), - (ADMIN, {}, "active", 409), - (ADMIN, {"direction": "reverse", "baseline_model": "my-router"}, None, 400), - (ADMIN, {"direction": "reverse", "baseline_model": "not/a real model!"}, None, 400), - (ADMIN, {"direction": "reverse", "baseline_model": "openai/gpt-4o", "router_name": "not-a-router"}, None, 400), + (NON_ADMIN, {}, (), 403), + (VIEWER, {}, (), 403), + (ADMIN, {"router_name": "not-a-router"}, (), 400), + (ADMIN, {"judge_model": "not/a real model!"}, (), 400), + (ADMIN, {"judge_model": "my-router"}, (), 400), + (ADMIN, {}, ("key-hash",), 409), + (ADMIN, {"api_key_ids": ("key-hash", "key-hash-2")}, ("key-hash-2",), 409), + (ADMIN, {"direction": "reverse", "baseline_model": "my-router"}, (), 400), + (ADMIN, {"direction": "reverse", "baseline_model": "not/a real model!"}, (), 400), + (ADMIN, {"direction": "reverse", "baseline_model": "openai/gpt-4o", "router_name": "not-a-router"}, (), 400), ], ids=[ "non-admin", @@ -621,23 +702,143 @@ async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones "unresolvable-judge", "router-as-judge", "already-active", + "one-of-several-keys-already-active", "router-as-baseline", "unresolvable-baseline", "reverse-still-needs-an-auto-router", ], ) async def test_start_shadow_eval_rejections( - monkeypatch: pytest.MonkeyPatch, caller, request_overrides, active, expected_status + monkeypatch: pytest.MonkeyPatch, caller, request_overrides, claimed, expected_status ): import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma(active_job=_job_record() if active else None) + prisma = _shadow_prisma(legs=[_leg_record(id=f"leg-{key}", group_id="job-7", api_key_id=key) for key in claimed]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) with pytest.raises(HTTPException) as exc: await start_shadow_eval(_start_request(**request_overrides), caller) assert exc.value.status_code == expected_status + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_names_the_busy_key_and_its_job(monkeypatch: pytest.MonkeyPatch): + """A key busy elsewhere blocks the whole start rather than being silently dropped from + it, and the 409 names which key and which job so the caller can stop or drop it.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(id="leg-b", group_id="job-7", api_key_id="key-hash-2")]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) + assert exc.value.status_code == 409 + assert "key-hash-2 (job job-7)" in exc.value.detail + + +@pytest.mark.asyncio +async def test_start_shadow_eval_reuses_a_key_whose_previous_job_already_stopped(monkeypatch: pytest.MonkeyPatch): + """The claim is held by unstopped legs only, matching the partial unique index. A read + that forgets that would strand every key that has ever finished a job.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(group_id="job-7", stopped_at=datetime.now(timezone.utc))]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + job = await start_shadow_eval(_start_request(), ADMIN) + + assert job.status == "running" + prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot(monkeypatch: pytest.MonkeyPatch): + """The two directions ask opposite questions of the same key, so a forward job holding + the slot must not block a reverse one. The second reverse start still 409s.""" + import litellm.proxy.proxy_server as proxy_server + + legs = [_leg_record(group_id="job-fwd")] + prisma = _shadow_prisma(legs=legs) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + reverse = _start_request(direction="reverse", baseline_model="openai/gpt-4o") + response = await start_shadow_eval(reverse, ADMIN) + + assert (response.direction, response.baseline_model) == ("reverse", "openai/gpt-4o") + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert rows[0]["direction"] == "reverse" + assert rows[0]["baseline_model"] == "openai/gpt-4o" + + legs.append(_leg_record(id="leg-2", group_id="job-rev", direction="reverse")) + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(reverse, ADMIN) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_start_shadow_eval_forward_leaves_the_baseline_column_empty(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + await start_shadow_eval(_start_request(), ADMIN) + + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert rows[0]["direction"] == "forward" + assert rows[0]["baseline_model"] is None + + +@pytest.mark.asyncio +async def test_start_shadow_eval_rejects_keys_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch): + """A typo'd api_key_id would otherwise create a leg no traffic can ever match. Every + unknown key is named at once, so a caller passing several fixes them in one round.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(known_keys=("key-hash",)) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(api_key_ids=("key-hash", "typo-a", "typo-b")), ADMIN) + assert exc.value.status_code == 400 + assert "typo-a, typo-b" in exc.value.detail + assert "key-hash," not in exc.value.detail + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +def test_start_shadow_eval_request_dedupes_and_bounds_the_key_set(): + """A key named twice would collide with itself on the one-active-per-key index, a job + scoping no key samples nothing, and the key-count cap bounds every downstream read.""" + assert _start_request(api_key_ids=("a", "b", "a")).api_key_ids == ("a", "b") + assert len(_start_request(api_key_ids=tuple(f"k{i}" for i in range(100))).api_key_ids) == 100 + with pytest.raises(ValidationError): + _start_request(api_key_ids=()) + with pytest.raises(ValidationError): + _start_request(api_key_ids=tuple(f"k{i}" for i in range(101))) + + +@pytest.mark.asyncio +async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + from prisma.errors import UniqueViolationError + + prisma = _shadow_prisma() + prisma.db.litellm_shadowevaljob.create_many = AsyncMock( + side_effect=UniqueViolationError(MagicMock(message="unique constraint")) + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(), ADMIN) + assert exc.value.status_code == 409 @pytest.mark.parametrize( @@ -657,97 +858,25 @@ def test_start_request_pins_baseline_model_to_reverse(overrides): @pytest.mark.asyncio -async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot(monkeypatch: pytest.MonkeyPatch): - """The two directions ask opposite questions of the same key, so a forward job holding - the slot must not block a reverse one. The second reverse start still 409s.""" - import litellm.proxy.proxy_server as proxy_server - - prisma = _shadow_prisma() - active = {"forward": _job_record()} - prisma.db.litellm_shadowevaljob.find_first = AsyncMock( - side_effect=lambda where, **_: active.get(str(where.get("direction"))) - ) - prisma.db.litellm_shadowevaljob.create = AsyncMock( - return_value=_job_record(direction="reverse", baseline_model="openai/gpt-4o") - ) - monkeypatch.setattr(proxy_server, "prisma_client", prisma) - monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - - reverse = _start_request(direction="reverse", baseline_model="openai/gpt-4o") - response = await start_shadow_eval(reverse, ADMIN) - - assert (response.direction, response.baseline_model) == ("reverse", "openai/gpt-4o") - create_data = prisma.db.litellm_shadowevaljob.create.call_args.kwargs["data"] - assert create_data["direction"] == "reverse" - assert create_data["baseline_model"] == "openai/gpt-4o" - - active["reverse"] = _job_record(id="job-2", direction="reverse") - with pytest.raises(HTTPException) as exc: - await start_shadow_eval(reverse, ADMIN) - assert exc.value.status_code == 409 - - -@pytest.mark.asyncio -async def test_start_shadow_eval_forward_leaves_the_baseline_column_empty(monkeypatch: pytest.MonkeyPatch): - import litellm.proxy.proxy_server as proxy_server - - prisma = _shadow_prisma() - monkeypatch.setattr(proxy_server, "prisma_client", prisma) - monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - - await start_shadow_eval(_start_request(), ADMIN) - - create_data = prisma.db.litellm_shadowevaljob.create.call_args.kwargs["data"] - assert create_data["direction"] == "forward" - assert create_data["baseline_model"] is None - - -@pytest.mark.asyncio -async def test_start_shadow_eval_rejects_a_key_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch): - """A typo'd api_key_id would otherwise create a job no traffic can ever match.""" - import litellm.proxy.proxy_server as proxy_server - - prisma = _shadow_prisma() - prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) - monkeypatch.setattr(proxy_server, "prisma_client", prisma) - monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - - with pytest.raises(HTTPException) as exc: - await start_shadow_eval(_start_request(), ADMIN) - assert exc.value.status_code == 400 - assert "not a key on this proxy" in exc.value.detail - - -@pytest.mark.asyncio -async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch): - import litellm.proxy.proxy_server as proxy_server - from prisma.errors import UniqueViolationError - - prisma = _shadow_prisma() - prisma.db.litellm_shadowevaljob.create = AsyncMock( - side_effect=UniqueViolationError(MagicMock(message="unique constraint")) - ) - monkeypatch.setattr(proxy_server, "prisma_client", prisma) - monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - - with pytest.raises(HTTPException) as exc: - await start_shadow_eval(_start_request(), ADMIN) - assert exc.value.status_code == 409 - - -@pytest.mark.asyncio -async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(monkeypatch: pytest.MonkeyPatch): +async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monkeypatch: pytest.MonkeyPatch): + """One read answers for every leg: totals and stratifications aggregate over the + group's leg ids, and the by-key slice maps each leg id back to its key hash.""" import litellm.proxy.proxy_server as proxy_server tier_rows = [ {"grp": "SIMPLE", "turn_count": 8, "real_wins": 2, "shadow_wins": 4, "ties": 2, "avg_confidence": 0.8}, {"grp": "REASONING", "turn_count": 2, "real_wins": 2, "shadow_wins": 0, "ties": 0, "avg_confidence": 0.9}, ] - prisma = _shadow_prisma(agg_rows=tier_rows) - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record()) - prisma.db.litellm_shadowevalattempt.find_first = AsyncMock( - return_value=MagicMock(error="judge call failed: boom") + leg_rows = [ + {"grp": "leg-1", "turn_count": 6, "real_wins": 1, "shadow_wins": 4, "ties": 1, "avg_confidence": 0.7}, + {"grp": "leg-2", "turn_count": 4, "real_wins": 3, "shadow_wins": 0, "ties": 1, "avg_confidence": 0.6}, + ] + prisma = _shadow_prisma( + legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=50)], + agg_rows=tier_rows, + by_leg_rows=leg_rows, ) + prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=MagicMock(error="judge call failed: boom")) monkeypatch.setattr(proxy_server, "prisma_client", prisma) response = await get_shadow_eval_job("job-1", VIEWER) @@ -762,6 +891,13 @@ async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(m assert response.results.by_tier[0].shadow_win_rate_pct == 50.0 assert response.results.overall_shadow_win_rate_pct == 40.0 assert response.results.overall_tie_rate_pct == 20.0 + assert [(s.group, s.turn_count) for s in response.results.by_key] == [("key-hash", 6), ("key-hash-2", 4)] + assert response.results.by_key[0].shadow_win_rate_pct == 66.7 + assert [(key.api_key_id, key.max_turns) for key in response.keys] == [("key-hash", 200), ("key-hash-2", 50)] + totals_args = [call.args for call in prisma.db.query_raw.await_args_list if "judged_count" in call.args[0]] + assert totals_args == [(totals_args[0][0], ["leg-1", "leg-2"])] + error_where = prisma.db.litellm_shadowevalattempt.find_first.call_args.kwargs["where"] + assert error_where == {"job_id": {"in": ["leg-1", "leg-2"]}, "outcome": "error"} @pytest.mark.asyncio @@ -780,79 +916,326 @@ async def test_get_shadow_eval_job_404s_and_gates_on_role(monkeypatch: pytest.Mo @pytest.mark.asyncio -async def test_list_shadow_eval_jobs_returns_derived_status_without_aggregates(monkeypatch: pytest.MonkeyPatch): +async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monkeypatch: pytest.MonkeyPatch): + """A job over two keys is one list entry with both keys, not two entries, and a job + whose keys all stopped reads stopped while a half-stopped one still runs.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma() - prisma.db.litellm_shadowevaljob.find_many = AsyncMock( - return_value=[ - _job_record(), - _job_record(id="job-2", ends_at=datetime.now(timezone.utc) - timedelta(days=1)), - _job_record(id="job-3", stopped_at=datetime.now(timezone.utc)), + stamp = datetime.now(timezone.utc) + prisma = _shadow_prisma( + legs=[ + _leg_record(created_at=datetime(2026, 8, 13, tzinfo=timezone.utc)), + _leg_record( + id="leg-2", + api_key_id="key-hash-2", + stopped_at=stamp, + created_at=datetime(2026, 8, 13, tzinfo=timezone.utc), + ), + _leg_record( + id="leg-3", + group_id="job-2", + stopped_at=stamp, + created_at=datetime(2026, 8, 12, tzinfo=timezone.utc), + ), + _leg_record( + id="leg-4", + group_id="job-3", + ends_at=datetime.now(timezone.utc) - timedelta(days=1), + created_at=datetime(2026, 8, 11, tzinfo=timezone.utc), + ), ] ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) - assert [job.status for job in jobs] == ["running", "completed", "stopped"] - swept = ShadowEvalJobResponse.model_validate( - _job_record( - id="job-4", - ends_at=datetime.now(timezone.utc) - timedelta(days=1), - stopped_at=datetime.now(timezone.utc), - ), - from_attributes=True, - ) - assert swept.status == "completed" + assert [(job.job_id, job.status) for job in jobs] == [ + ("job-1", "running"), + ("job-2", "stopped"), + ("job-3", "completed"), + ] + assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"] assert all(job.judged_count is None and job.results is None for job in jobs) - assert prisma.db.query_raw.await_count == 0 + legs_sql, legs_limit = prisma.db.query_raw.await_args_list[0].args + assert "GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int" in legs_sql + assert legs_limit == 50 + counts_sql, _ = prisma.db.query_raw.await_args_list[1].args + assert "AS attempt_count" in counts_sql + assert "j.stopped_at IS NULL OR a.created_at <= j.stopped_at" in counts_sql + assert prisma.db.query_raw.await_count == 2 + prisma.db.litellm_shadowevaljob.find_many.assert_not_called() @pytest.mark.asyncio -async def test_shadow_eval_responses_name_the_shadowed_key(monkeypatch: pytest.MonkeyPatch): +async def test_list_shadow_eval_jobs_filters_to_jobs_containing_the_key(monkeypatch: pytest.MonkeyPatch): + """The filter matches a key anywhere in a job's key set and still returns the whole + job, sibling keys included.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma() - prisma.db.litellm_shadowevaljob.find_many = AsyncMock( - return_value=[_job_record(), _job_record(id="job-2", api_key_id="deleted-key-hash")] + prisma = _shadow_prisma( + legs=[ + _leg_record(), + _leg_record(id="leg-2", api_key_id="key-hash-2"), + _leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash-2"), + _leg_record(id="leg-4", group_id="job-3"), + ] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id="key-hash-2", limit=50) + + assert [job.job_id for job in jobs] == ["job-1", "job-2"] + assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"] + + +@pytest.mark.parametrize( + ("stopped_flags", "days_left", "expected"), + [ + ((False, False), 7, "running"), + ((True, False), 7, "running"), + ((True, True), 7, "stopped"), + ((True, True), -1, "completed"), + ((False, False), -1, "completed"), + ], +) +@pytest.mark.asyncio +async def test_job_status_runs_until_every_key_stops_and_completed_outranks_stopped( + monkeypatch: pytest.MonkeyPatch, stopped_flags: tuple[bool, ...], days_left: int, expected: str +): + import litellm.proxy.proxy_server as proxy_server + + stamp = datetime.now(timezone.utc) + prisma = _shadow_prisma( + legs=[ + _leg_record( + id=f"leg-{index}", + api_key_id=f"key-{index}", + stopped_at=stamp if stopped else None, + ends_at=datetime.now(timezone.utc) + timedelta(days=days_left), + ) + for index, stopped in enumerate(stopped_flags) + ] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + + assert [job.status for job in jobs] == [expected] + + +@pytest.mark.asyncio +async def test_list_reads_completed_once_every_key_spends_its_budget(monkeypatch: pytest.MonkeyPatch): + """A job whose keys all exhausted their turn budgets stopped sampling on its own, so + it must read completed on the very next list, before any sweep stamps its legs; one + key under budget keeps the whole job running. An operator starting an unrelated eval + must never look like it terminated a finished one.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma( + legs=[ + _leg_record(max_turns=5), + _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=5), + _leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash", max_turns=5), + _leg_record(id="leg-4", group_id="job-2", api_key_id="key-hash-2", max_turns=5), + ] + ) + prisma.attempt_rows = [ + {"job_id": "leg-1", "attempt_count": 5}, + {"job_id": "leg-2", "attempt_count": 6}, + {"job_id": "leg-3", "attempt_count": 5}, + {"job_id": "leg-4", "attempt_count": 3}, + ] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + + by_id = {job.job_id: job for job in jobs} + assert by_id["job-1"].status == "completed" + assert all(key.stopped_at is None for key in by_id["job-1"].keys) + assert by_id["job-2"].status == "running" + assert {key.api_key_id: key.attempt_count for key in by_id["job-2"].keys} == {"key-hash": 5, "key-hash-2": 3} + + +@pytest.mark.asyncio +async def test_recorded_operator_stop_outranks_budget_arithmetic(monkeypatch: pytest.MonkeyPatch): + """A detached attempt can land around the stop and push the raw count past the + budget; the recorded stopped_by must keep the job reading stopped regardless.""" + import litellm.proxy.proxy_server as proxy_server + + stamp = datetime.now(timezone.utc) + prisma = _shadow_prisma(legs=[_leg_record(max_turns=5, stopped_at=stamp, stopped_by="admin")]) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + assert jobs[0].status == "stopped" + assert jobs[0].stopped_by == "admin" + + detail = await get_shadow_eval_job("job-1", VIEWER) + assert detail.status == "stopped" + + +@pytest.mark.asyncio +async def test_backfilled_legacy_stop_never_reads_as_completion(monkeypatch: pytest.MonkeyPatch): + """Jobs stopped before stopped_by existed are backfilled with 'unknown' by the + migration, so even one whose stray attempts crossed the budget stays stopped.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma( + legs=[_leg_record(max_turns=5, stopped_at=datetime.now(timezone.utc), stopped_by="unknown")] + ) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + assert jobs[0].status == "stopped" + + +def test_stopped_by_migration_backfills_every_job_that_displayed_stopped(): + """The migration must close the pre-column population: without the backfill, a + legacy stop whose stray attempts crossed the budget would read completed.""" + import litellm_proxy_extras + + sql = ( + Path(litellm_proxy_extras.__file__).parent + / "migrations" + / "20260818224500_add_shadow_eval_stopped_by" + / "migration.sql" + ).read_text() + assert 'ADD COLUMN "stopped_by" TEXT' in sql + assert "SET stopped_by = 'unknown'" in sql + assert "WHERE stopped_at IS NOT NULL AND ends_at > (NOW() AT TIME ZONE 'utc')" in sql + + +@pytest.mark.asyncio +async def test_stop_rejects_a_job_that_already_spent_its_budget(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(max_turns=3)]) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 3}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with pytest.raises(HTTPException) as exhausted: + await stop_shadow_eval_job("job-1", ADMIN) + assert exhausted.value.status_code == 400 + assert "completed" in exhausted.value.detail + prisma.db.litellm_shadowevaljob.update_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_shadow_eval_responses_name_every_shadowed_key(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma( + legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="deleted-key-hash")], + known_keys=("key-hash", "key-hash-2"), ) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - started = await start_shadow_eval(_start_request(), ADMIN) - assert (started.key_alias, started.key_name) == ("prod-alpha", "sk-...lpha") - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) - assert [(job.key_alias, job.key_name) for job in jobs] == [("prod-alpha", "sk-...lpha"), (None, None)] + assert [(key.key_alias, key.key_name) for key in jobs[0].keys] == [ + (None, None), + ("prod-alpha", "sk-...lpha"), + ] batched_where = prisma.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] assert batched_where == {"token": {"in": ["deleted-key-hash", "key-hash"]}} - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record()) detail = await get_shadow_eval_job("job-1", VIEWER) - assert detail.key_alias == "prod-alpha" + assert [key.key_alias for key in detail.keys] == [None, "prod-alpha"] @pytest.mark.asyncio -async def test_stop_shadow_eval_sets_stopped_at_and_rejects_non_running(monkeypatch: pytest.MonkeyPatch): +async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_running( + monkeypatch: pytest.MonkeyPatch, +): + """One stop ends sampling for the whole job, while a leg that already stopped on its + own budget keeps the stopped_at it earned.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma() - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record()) + earned = datetime.now(timezone.utc) - timedelta(hours=1) + prisma = _shadow_prisma(legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", stopped_at=earned)]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) stopped = await stop_shadow_eval_job("job-1", ADMIN) - assert stopped.status == "stopped" - update = prisma.db.litellm_shadowevaljob.update.call_args.kwargs - assert set(update["data"]) == {"stopped_at"} - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock( - return_value=_job_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1)) - ) + assert stopped.status == "stopped" + assert stopped.stopped_by == "admin" + stop_sql, stop_group, stop_operator, stop_stamp = prisma.db.execute_raw.call_args.args + assert "SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp)" in stop_sql + assert "WHERE group_id = $1 AND stopped_by IS NULL" in stop_sql + assert "ends_at > (NOW() AT TIME ZONE 'utc')" in stop_sql + assert ") < k.max_turns" in stop_sql + assert (stop_group, stop_operator) == ("job-1", "admin") + assert datetime.fromisoformat(stop_stamp).tzinfo is None + assert prisma.db.execute_raw.await_count == 1 + prisma.db.litellm_shadowevaljob.update_many.assert_not_called() + by_key = {key.api_key_id: key.stopped_at for key in stopped.keys} + assert by_key["key-hash-2"] == earned + assert by_key["key-hash"] is not None and by_key["key-hash"] != earned + + done_leg = _leg_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1)) + prisma_done = _shadow_prisma(legs=[done_leg]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_done) with pytest.raises(HTTPException) as exc: await stop_shadow_eval_job("job-1", ADMIN) assert exc.value.status_code == 400 + assert "already completed" in exc.value.detail + assert done_leg.stopped_by is None with pytest.raises(HTTPException) as forbidden: await stop_shadow_eval_job("job-1", VIEWER) assert forbidden.value.status_code == 403 + + +def test_every_shadow_eval_sql_constant_speaks_naive_utc(): + """The tables store naive UTC wall time (prisma's convention), so SQL-side time must be + NOW() AT TIME ZONE 'utc' and python-side params must cast ::timestamp; a bare NOW() or a + timestamptz cast writes session-local wall time into the naive column and skews every + comparison against prisma-written stamps.""" + import litellm.proxy.management_endpoints.auto_router_endpoints as module + + sql_constants = {name: value for name, value in vars(module).items() if name.endswith("_SQL")} + assert sql_constants + for name, sql in sql_constants.items(): + assert "::timestamptz" not in sql, name + for occurrence in sql.split("NOW()")[1:]: + assert occurrence.startswith(" AT TIME ZONE 'utc'"), name + + +@pytest.mark.asyncio +async def test_a_stop_racing_the_last_budgeted_attempt_reports_completed_not_stopped( + monkeypatch: pytest.MonkeyPatch, +): + """The statement claims the job only while a leg still samples, so a stop landing in + the same instant the budget spends records nothing and the job keeps reading + completed; stamping it would misreport a self-ended job as operator-stopped forever.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(max_turns=2)]) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 2}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with pytest.raises(HTTPException) as exc: + await stop_shadow_eval_job("job-1", ADMIN) + assert exc.value.status_code == 400 + assert "already completed" in exc.value.detail + assert prisma.db.litellm_shadowevaljob.find_many.await_args.kwargs["where"] == {"group_id": "job-1"} + + +@pytest.mark.asyncio +async def test_two_racing_stops_produce_exactly_one_winner(monkeypatch: pytest.MonkeyPatch): + """The statement's stopped_by IS NULL predicate lets only one racer claim rows; the + loser reads the stamped state and gets the same answer a late caller gets.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record()]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + first = await stop_shadow_eval_job("job-1", ADMIN) + assert first.status == "stopped" + + with pytest.raises(HTTPException) as exc: + await stop_shadow_eval_job("job-1", ADMIN) + assert exc.value.status_code == 400 + assert "already stopped" in exc.value.detail diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 7bb550f729e..b307e3d0f2a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -77,7 +77,15 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ baseline_model: null, judge_model: "anthropic/claude-sonnet-5", shadow_percentage: 10, - max_turns: 200, + keys: [ + { + api_key_id: "hashed-key-abc", + max_turns: 200, + stopped_at: null, + key_alias: "prod-alpha", + key_name: "sk-...alpha", + }, + ], judged_count: 42, error_count: 1, judge_spend: 3.21, @@ -110,19 +118,28 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ avg_judge_confidence: 0.8, }, ], + by_key: [], overall_shadow_win_rate_pct: 48.0, overall_tie_rate_pct: 22.0, }, created_at: "2026-08-07T00:00:00Z", ends_at: "2026-09-07T00:00:00Z", - stopped_at: null, - api_key_id: "hashed-key-abc", - key_alias: "prod-alpha", - key_name: "sk-...alpha", last_error: null, ...overrides, }); +const keyEntry = ( + api_key_id: string, + overrides: Partial = {}, +): ShadowEvalJob["keys"][number] => ({ + api_key_id, + max_turns: 200, + stopped_at: null, + key_alias: null, + key_name: null, + ...overrides, +}); + const mockHooks = ({ jobs = [], detailsById = {}, @@ -199,8 +216,8 @@ describe("ShadowEvalSection", () => { it("gives every active job its own card with a stop button, with the form still offered", () => { mockHooks({ jobs: [ - job({ job_id: "job-a", status: "running", api_key_id: "key-a" }), - job({ job_id: "job-b", status: "running", api_key_id: "key-b" }), + job({ job_id: "job-a", status: "running", keys: [keyEntry("key-a")] }), + job({ job_id: "job-b", status: "running", keys: [keyEntry("key-b")] }), ], }); render(); @@ -342,7 +359,7 @@ describe("ShadowEvalSection", () => { expect(container).toBeEmptyDOMElement(); }); - it("keeps the start button disabled until key, router, and judge model are picked, then submits them", async () => { + it("keeps the start button disabled until key, router, and judge model are picked, then submits the key as a list", async () => { const user = userEvent.setup(); const { start } = mockHooks({}); render(); @@ -361,7 +378,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByText("Start shadow eval")); const expectedBody = { - api_key_id: "hash-alpha", + api_key_ids: ["hash-alpha"], router_name: "gpt-auto", direction: "forward", shadow_percentage: 10, @@ -396,7 +413,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByText("Start shadow eval")); const expectedBody = { - api_key_id: "hash-alpha", + api_key_ids: ["hash-alpha"], router_name: "gpt-auto", direction: "reverse", baseline_model: "prod-claude", @@ -429,9 +446,9 @@ describe("ShadowEvalSection", () => { }); it("labels the shadowed key by alias, then masked name, then truncated hash", () => { - expect(shadowedKeyLabel(job())).toBe("prod-alpha"); - expect(shadowedKeyLabel(job({ key_alias: null }))).toBe("sk-...alpha"); - expect(shadowedKeyLabel(job({ key_alias: null, key_name: null }))).toBe("hashed-key…"); + expect(shadowedKeyLabel(job().keys[0])).toBe("prod-alpha"); + expect(shadowedKeyLabel(keyEntry("hashed-key-abc", { key_name: "sk-...alpha" }))).toBe("sk-...alpha"); + expect(shadowedKeyLabel(keyEntry("hashed-key-abc"))).toBe("hashed-key…"); }); it("keeps an older job's verdicts reachable through the previous evaluations list", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index 005636615b1..6d240a84c98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -24,6 +24,7 @@ import { useStartShadowEval, useStopShadowEval, type ShadowEvalJob, + type ShadowEvalJobKey, type ShadowEvalSlice, } from "./useShadowEval"; @@ -50,19 +51,24 @@ const routerMatchedOrBeatPct = ( ? 100 - results.overall_shadow_win_rate_pct : results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct; -export const shadowedKeyLabel = (job: ShadowEvalJob): string => - job.key_alias || job.key_name || `${job.api_key_id.slice(0, 10)}…`; +export const shadowedKeyLabel = (key: ShadowEvalJobKey): string => + key.key_alias || key.key_name || `${key.api_key_id.slice(0, 10)}…`; + +const shadowedKeysLabel = (job: ShadowEvalJob): string => + job.keys.length === 1 ? shadowedKeyLabel(job.keys[0]) : `${job.keys.length} keys`; + +const totalBudget = (job: ShadowEvalJob): number => job.keys.reduce((sum, key) => sum + key.max_turns, 0); const jobHeadline = (job: ShadowEvalJob): React.ReactNode => job.direction === "reverse" ? ( <> Comparing {job.router_name} to{" "} {job.baseline_model} on {job.shadow_percentage}% of{" "} - {shadowedKeyLabel(job)} traffic + {shadowedKeysLabel(job)} traffic ) : ( <> - Shadowing {job.shadow_percentage}% of {shadowedKeyLabel(job)} traffic + Shadowing {job.shadow_percentage}% of {shadowedKeysLabel(job)} traffic via {job.router_name} ); @@ -178,7 +184,8 @@ const emptyResultsText = (job: ShadowEvalJob, resultsError: boolean): string => const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({ job, resultsError = false }) => { const results = job.results; - if (!results || (results.by_tier.length === 0 && results.by_current_model.length === 0)) { + const stratifications = results ? [results.by_tier, results.by_current_model, results.by_key] : []; + if (!results || stratifications.every((slices) => slices.length === 0)) { return

{emptyResultsText(job, resultsError)}

; } return ( @@ -224,7 +231,7 @@ const JobResults: React.FC<{

{jobHeadline(job)}

- {(job.judged_count ?? 0).toLocaleString()} of {job.max_turns.toLocaleString()} turns judged ·{" "} + {(job.judged_count ?? 0).toLocaleString()} of {totalBudget(job).toLocaleString()} turns judged ·{" "} {(job.error_count ?? 0).toLocaleString()} errored · {usd(job.judge_spend ?? 0)} judge spend {active && remaining ? ` · ${remaining}` : ""}

@@ -386,14 +393,13 @@ const StartForm: React.FC = () => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxTurns = Number.parseInt(maxTurns, 10); const maxTurnsValid = parsedMaxTurns >= 1 && parsedMaxTurns <= 2000; - const filled = - [apiKeyId, routerName, judgeModel].every((field) => field !== "") && - (direction === "forward" || baselineModel !== ""); + const baselinePicked = direction === "forward" || baselineModel !== ""; + const filled = [apiKeyId, routerName, judgeModel].every((field) => field !== "") && baselinePicked; const boundsValid = percentageValid && maxTurnsValid; const valid = Boolean(accessToken) && filled && boundsValid; const handleStart = () => { const startBody = { - api_key_id: apiKeyId, + api_key_ids: [apiKeyId], router_name: routerName, direction, ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts index 7645fcc3346..eef98320e67 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts @@ -7,6 +7,7 @@ import { $api, fetchClient } from "@/lib/http/api"; import type { components } from "@/lib/http/schema"; export type ShadowEvalJob = components["schemas"]["ShadowEvalJobResponse"]; +export type ShadowEvalJobKey = components["schemas"]["ShadowEvalJobKeyResponse"]; export type ShadowEvalSlice = components["schemas"]["ShadowEvalSlice"]; export type StartShadowEvalRequest = components["schemas"]["StartShadowEvalRequest"]; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5794c60e97b..4c63586d4d6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -816,7 +816,8 @@ export interface paths { }; /** * List Shadow Eval Jobs - * @description List shadow eval jobs, newest first. Counts and results ride the detail endpoint only. + * @description List shadow eval jobs, newest first, each key with its attempt count so status is + * accurate. Judged counts, spend, and results ride the detail endpoint only. */ get: operations["list_shadow_eval_jobs_auto_router_shadow_eval_get"]; put?: never; @@ -838,20 +839,21 @@ export interface paths { put?: never; /** * Start Shadow Eval - * @description Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second - * arm, judge the two responses blind, and stratify win rates by tier and by the model that - * served the real arm. + * @description Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against + * a second arm, judge the two responses blind, and stratify win rates by tier, by the model + * that served the real arm, and by key. * - * A forward job answers whether the key should adopt router_name: it samples the requests + * A forward job answers whether the keys should adopt router_name: it samples the requests * the router did not serve and duplicates them through it. A reverse job answers whether a * key already on the router still gains from it: it samples the requests the router did * serve and duplicates them against baseline_model. A key can hold one active job per * direction, so both questions can run at once. * - * Shadow responses are never served to users. The job samples until it has judged - * max_turns turns, reaches the end of its window, or is stopped; sampling changes - * propagate to pods within about 10 seconds. Shadow and judge calls bill to the - * shadowed key but are excluded from request counts and auto-router adoption metrics. + * Shadow responses are never served to users. Each key samples until it has judged + * max_turns turns of its own traffic, the job's window ends, or the job is stopped, so one + * key running out of budget does not end sampling for the others; sampling changes + * propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed + * key but are excluded from request counts and auto-router adoption metrics. */ post: operations["start_shadow_eval_auto_router_shadow_eval_start_post"]; delete?: never; @@ -891,7 +893,12 @@ export interface paths { put?: never; /** * Stop Shadow Eval Job - * @description Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s. + * @description Stop an active shadow eval job, every key it scopes at once. Attempts are kept; + * sampling halts within ~10s. Keys that already stopped on their own budget keep the + * stopped_at they earned. The statement is the whole state machine: it claims the job + * only while a leg still samples inside the window with no stop recorded, so a racing + * operator, a same-instant budget spend, and a repeat stop all read the same 400 with + * the status the job actually holds. */ post: operations["stop_shadow_eval_job_auto_router_shadow_eval__job_id__stop_post"]; delete?: never; @@ -33201,18 +33208,49 @@ export interface components { timeout?: number | null; }; /** - * ShadowEvalJobResponse - * @description A shadow-eval job. Validates directly from the prisma record (job_id reads the - * row's id); status is derived from stopped_at and ends_at, never stored, so no writer - * anywhere can produce an inconsistent one. Aggregate fields are populated by the - * detail endpoint only and stay None on list responses. + * ShadowEvalJobKeyResponse + * @description One key a job shadows, with its own budget and stop state. */ - ShadowEvalJobResponse: { + ShadowEvalJobKeyResponse: { /** * Api Key Id - * @description The hashed virtual key whose traffic this job evaluates, and only that key's + * @description The hashed virtual key whose traffic this entry scopes */ api_key_id: string; + /** + * Attempt Count + * @description This key's sampled attempts so far, judged and errored alike, the same count the sampler budgets against max_turns; populated on list and detail responses. Frozen at stopped_at once the key is stamped, so in-flight attempts landing after a stop never reclassify it + */ + attempt_count?: number | null; + /** + * Key Alias + * @description Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted + */ + key_alias?: string | null; + /** + * Key Name + * @description Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias + */ + key_name?: string | null; + /** + * Max Turns + * @description This key's own sample budget, independent of its siblings' + */ + max_turns: number; + /** + * Stopped At + * @description When this key's slot was stamped free, whether its own budget ran out, the window closed, or an operator stopped the job; status is derived, so a spent budget reads completed even while this is still unset + */ + stopped_at?: string | null; + }; + /** + * ShadowEvalJobResponse + * @description A shadow-eval job over one or more keys, each with its own budget and stop state; + * status is derived from stopped_by, the keys' stop and budget state, and ends_at, + * never stored, so no writer anywhere can produce an inconsistent one. Aggregate + * fields are populated by the detail endpoint only and stay None on list responses. + */ + ShadowEvalJobResponse: { /** Baseline Model */ baseline_model?: string | null; /** @@ -33251,22 +33289,15 @@ export interface components { */ judged_count?: number | null; /** - * Key Alias - * @description Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted + * Keys + * @description The keys whose traffic this job evaluates, and only those keys', each with its own budget */ - key_alias?: string | null; - /** - * Key Name - * @description Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias - */ - key_name?: string | null; + keys: components["schemas"]["ShadowEvalJobKeyResponse"][]; /** * Last Error * @description Most recent attempt error; detail endpoint only */ last_error?: string | null; - /** Max Turns */ - max_turns: number; /** @description Stratified verdicts; detail endpoint only */ results?: components["schemas"]["ShadowEvalResult"] | null; /** Router Name */ @@ -33275,13 +33306,19 @@ export interface components { shadow_percentage: number; /** * Status - * @description A job whose window has passed reads completed even if a later sweep stamped - * stopped_at; stopped means sampling ended before the window did. + * @description Three recorded facts, no history-guessing: a stop is stopped_by (the migration + * backfills it for every job that displayed stopped when the column arrived, so the + * pre-column population is closed), completion is the window passing or every key + * spending its budget, and anything else is running. The all-keys-stamped fallback + * covers only stops written by pre-column pods during a rolling deploy. * @enum {string} */ readonly status: "running" | "completed" | "stopped"; - /** Stopped At */ - stopped_at?: string | null; + /** + * Stopped By + * @description The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled by migration for jobs that displayed stopped when the column arrived; None when the job ended on its own. Its presence is what makes a job read stopped rather than completed + */ + stopped_by?: string | null; }; /** * ShadowEvalResult @@ -33290,9 +33327,14 @@ export interface components { ShadowEvalResult: { /** * By Current Model - * @description Sliced by the model that served the real arm: the key's incumbent models in forward mode, and in reverse the models the router itself picked + * @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked */ by_current_model: components["schemas"]["ShadowEvalSlice"][]; + /** + * By Key + * @description One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job scopes but has not judged a turn for yet are absent rather than reported as zero + */ + by_key: components["schemas"]["ShadowEvalSlice"][]; /** By Tier */ by_tier: components["schemas"]["ShadowEvalSlice"][]; /** Overall Shadow Win Rate Pct */ @@ -33500,14 +33542,14 @@ export interface components { }; /** * StartShadowEvalRequest - * @description Start duplicating a key's traffic for blind comparison against an auto-router. + * @description Start duplicating one or more keys' traffic for blind comparison against an auto-router. */ StartShadowEvalRequest: { /** - * Api Key Id - * @description The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this key's traffic; requests made with any other key are not sampled. + * Api Key Ids + * @description The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these keys' traffic; requests made with any other key are not sampled. Each key carries its own max_turns budget, so one key exhausting its budget leaves the others sampling. At most 100 keys per job, which also bounds every read the job's endpoints make. */ - api_key_id: string; + api_key_ids: string[]; /** * Baseline Model * @description Required when direction is reverse and rejected otherwise: the fixed model the router's own responses are judged against. Must be a plain model rather than another auto-router @@ -33534,7 +33576,7 @@ export interface components { judge_model: string; /** * Max Turns - * @description Sample budget: the job judges at most this many turns, then completes. This is also the spend bound; expected judge cost is roughly max_turns times one judge call + * @description Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, so a job over N keys judges at most N times max_turns turns. This is also the spend bound; expected judge cost is roughly that turn ceiling times one judge call * @default 200 */ max_turns: number; @@ -37931,7 +37973,7 @@ export interface operations { list_shadow_eval_jobs_auto_router_shadow_eval_get: { parameters: { query?: { - /** @description Filter to jobs shadowing this key */ + /** @description Filter to jobs that shadow this key, alone or alongside others */ api_key_id?: string | null; /** @description Newest jobs to return */ limit?: number; From 3d51eb378a027b17e4f123b3eb768b31150bcccb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 19 Aug 2026 14:07:50 -0700 Subject: [PATCH 29/88] refactor(ui): migrate the antd Alert call sites onto the shared Alert (#37513) Moves all 33 antd Alert usages across 20 dashboard files onto src/components/shared/Alert, following the composition the rest of the dashboard already uses: message becomes AlertTitle, description becomes AlertDescription, showIcon becomes a lucide icon child, and closable becomes an AlertAction ghost button. antd type="success" has no counterpart on the shared Alert, so the two success sites land on the default variant with a CircleCheck icon, which is what cloudzero_export_modal and CloudZeroIntegrationSettings already do for the same case. LoginPage's dismissible SSO notice moves into its own SsoEnabledNotice component in the same file: antd's closable carried its own dismiss state, and inlining it pushed LoginPageContent past the complexity budget. Four files lose their last antd symbol, so their no-restricted-imports suppressions are pruned by hand. antd import sites drop from 115 to 111 across 107 to 103 files, and the no-restricted-imports ratchet drops from 119 to 115 over 110 to 106 files. One test asserted antd's own ant-alert-info class; it is repointed to the shared Alert's text-info variant class, which keeps the same "info, not warning" check. Every other colocated test passes untouched. --- ui/litellm-dashboard/eslint-suppressions.json | 18 -- .../admin-panel/_components/AdminPanel.tsx | 18 +- .../_components/MCPPermissionManagement.tsx | 20 +- .../_components/UserEnvVarsModal.tsx | 14 +- .../mcp-servers/_components/mcp_connect.tsx | 50 ++-- .../_components/mcp_server_edit.tsx | 20 +- .../policies/_components/add_policy_form.tsx | 60 ++--- .../_components/policy_test_panel.tsx | 10 +- .../_components/CreateVectorStore.tsx | 38 +-- .../_components/S3VectorsConfig.tsx | 52 ++-- .../_components/VectorStoreForm.tsx | 246 +++++++++--------- .../src/app/login/LoginPage.tsx | 119 +++++---- .../src/app/onboarding/OnboardingFormBody.tsx | 24 +- .../src/components/CreateUserButton.tsx | 29 +-- .../MCPSemanticFilterSettings.tsx | 59 +++-- .../src/components/add_model/AddModelForm.tsx | 18 +- .../src/components/add_pass_through.tsx | 20 +- .../PassThroughGuardrailsSection.tsx | 40 ++- .../user_search_modal.test.tsx | 2 +- .../common_components/user_search_modal.tsx | 18 +- .../organisms/RegenerateKeyModal.tsx | 10 +- .../update_model_credentials_modal.tsx | 18 +- 22 files changed, 460 insertions(+), 443 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 98225de5d8b..09e602e269b 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1404,11 +1404,6 @@ "count": 2 } }, - "src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx": { "no-nested-ternary": { "count": 2 @@ -1453,9 +1448,6 @@ "local/no-complex-jsx-arrow": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -1470,11 +1462,6 @@ "count": 1 } }, - "src/app/onboarding/OnboardingFormBody.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/AIHub/ModelHubTable.test.tsx": { "max-params": { "count": 1 @@ -1998,11 +1985,6 @@ "count": 1 } }, - "src/components/common_components/PassThroughGuardrailsSection.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/RateLimitTypeFormItem.test.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index 48eaa62e37b..e9d7190abf5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -7,8 +7,8 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Alert as AntdAlert, Modal, Space, Tabs, Typography } from "antd"; -import { Info } from "lucide-react"; +import { Modal, Space, Tabs, Typography } from "antd"; +import { Info, TriangleAlert } from "lucide-react"; import React, { useEffect, useState } from "react"; import NewBadge from "@/components/common_components/NewBadge"; import { useBaseUrl } from "@/components/constants"; @@ -223,12 +223,14 @@ const AdminPanel: React.FC = ({ proxySettings }) => { <> ✨ Security Settings - + + + SSO Configuration Deprecated + + Editing SSO Settings on this page is deprecated and will be removed in a future version. Please use the + SSO Settings tab for SSO configuration. + +
= ({ )} {showInternalDelegatePkceWarning && ( - + + + Internal server with upstream OAuth delegation + + This MCP server is configured as internal-only but delegates auth to upstream. Anonymous users will be + able to reach the upstream OAuth2 /authorize flow without a LiteLLM session. Ensure your upstream + provider and network enforce access controls. + + )} = ({ server, open, acces
) : isError ? ( - + + + Failed to load env vars + ) : required.length === 0 ? ( - + + + No per-user fields configured for this server. + ) : ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx index bb739901faa..8aeb94404f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx @@ -1,10 +1,22 @@ /* eslint-disable react/no-unescaped-entities */ import React, { useState } from "react"; -import { Card, Typography, Space, Alert, Switch } from "antd"; +import { Card, Typography, Space, Switch } from "antd"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { CopyIcon, Code, Terminal, Globe, CheckIcon, ExternalLinkIcon, KeyIcon, ServerIcon, Zap } from "lucide-react"; +import { + CopyIcon, + Code, + Terminal, + Globe, + CheckIcon, + ExternalLinkIcon, + Info, + KeyIcon, + ServerIcon, + Zap, +} from "lucide-react"; import { getProxyBaseUrl } from "@/components/networking"; import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; @@ -69,25 +81,21 @@ const FeatureCard: React.FC = ({
{useServerHeader && ( - -

- Option 1: Get a specific server: "{serverName.replace(/\s+/g, "_")}" -

-

- Option 2: Get a group of MCPs: "dev-group" -

-

- You can also mix both: "Server1,dev-group" -

-
- } - /> + + + Two Options + +

+ Option 1: Get a specific server: "{serverName.replace(/\s+/g, "_")}" +

+

+ Option 2: Get a group of MCPs: "dev-group" +

+

+ You can also mix both: "Server1,dev-group" +

+ + )}
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index eb462f788a5..94b66f1bdfa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -1,5 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Select, Tooltip, Input, InputNumber, Alert } from "antd"; +import { Select, Tooltip, Input, InputNumber } from "antd"; +import { TriangleAlert } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { FormProvider, useForm } from "react-hook-form"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button } from "@/components/ui/button"; @@ -1041,13 +1043,15 @@ const MCPServerEdit: React.FC = ({ {!isStdioTransport && isOAuthAuthType && ( <> {!oauthFlowTypeValue && !isDelegateAuth && ( - + + + This server has no OAuth flow set + + Choose Machine-to-Machine (M2M) or Interactive (PKCE) so LiteLLM authenticates it the way you + intend, then save. Until it is set, LiteLLM falls back to interactive per-user auth and treats + a machine-to-machine credential shape conservatively. + + )} = ({ {selectedMode === "flow_builder" && ( + > + + You'll be redirected to the full-screen Flow Builder to design your policy logic visually. + + )}
@@ -457,35 +461,33 @@ const AddPolicyForm: React.FC = ({ {resolvedGuardrails.length > 0 && ( - - - These are the final guardrails that will be applied (including inheritance): - -
- {resolvedGuardrails.map((g) => ( - - {g} - - ))} -
+ + + Resolved Guardrails + + + These are the final guardrails that will be applied (including inheritance): + +
+ {resolvedGuardrails.map((g) => ( + + {g} + + ))}
- } - type="info" - showIcon - /> + + )} - + + + Model Scope + + By default, this policy will run on all models. You can optionally restrict it to specific models below. + +
Model Condition Type diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.tsx index 90642d23a46..9c234a7a70c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.tsx @@ -1,6 +1,8 @@ import React, { useState, useEffect } from "react"; import { useForm } from "react-hook-form"; -import { Alert, Empty } from "antd"; +import { Empty } from "antd"; +import { CircleAlert } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { resolvePoliciesCall, teamListCall, keyListCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { FieldGroup } from "@/components/shared/form/field"; @@ -325,7 +327,11 @@ const PolicyTestPanel: React.FC = ({ accessToken }) => { )} {hasSearched && !result && !isLoading && ( - + + + Error + Failed to resolve policies. Check the proxy logs. + )}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 93aaa90ef7c..271ef31a68e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -1,9 +1,10 @@ import React, { useState } from "react"; -import { Upload, Alert } from "antd"; +import { Upload } from "antd"; import { toast } from "@/lib/toast"; import { InboxOutlined } from "@ant-design/icons"; import type { UploadProps } from "antd"; -import { CircleHelp } from "lucide-react"; +import { CircleCheck, CircleHelp, X } from "lucide-react"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { ragIngestCall } from "@/components/networking"; import { DocumentUpload, RAGIngestResponse } from "@/components/vector_store_management/types"; import DocumentsTable from "./DocumentsTable"; @@ -359,22 +360,23 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu {/* Success Message */} {ingestResults.length > 0 && ( - -

- Vector Store ID: {ingestResults[0]?.vector_store_id} -

-

- Documents Ingested: {ingestResults.length} -

-
- } - type="success" - showIcon - closable - /> + + + Vector Store Created Successfully + +

+ Vector Store ID: {ingestResults[0]?.vector_store_id} +

+

+ Documents Ingested: {ingestResults.length} +

+ + + + + )}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx index 3e286ee9b92..cdc63e8b7ec 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Alert } from "antd"; -import { CircleHelp } from "lucide-react"; +import { CircleHelp, Info } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import { Field, FieldError, FieldLabel } from "@/components/shared/form/field"; import { @@ -72,32 +72,28 @@ const S3VectorsConfig: React.FC = ({ accessToken, provider return ( - -

AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

-
    -
  • Vector buckets and indexes will be automatically created if they don't exist
    • -
  • Vector dimensions are auto-detected from your selected embedding model
    • -
  • Ensure your AWS credentials have permissions for S3 Vectors operations
    • -
  • - Learn more:{" "} - - AWS S3 Vectors Documentation - -
    • -
-
- } - type="info" - showIcon - style={{ marginBottom: "16px" }} - /> + + + AWS S3 Vectors Setup + +

AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

+
    +
  • Vector buckets and indexes will be automatically created if they don't exist
    • +
  • Vector dimensions are auto-detected from your selected embedding model
    • +
  • Ensure your AWS credentials have permissions for S3 Vectors operations
    • +
  • + Learn more:{" "} + + AWS S3 Vectors Documentation + +
    • +
+ + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx index 05c829d79df..ff6c10bf428 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Modal, Alert } from "antd"; -import { CircleHelp, Eye, EyeOff } from "lucide-react"; +import { Modal } from "antd"; +import { CircleHelp, Eye, EyeOff, Info } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { useWatch } from "react-hook-form"; import { z } from "zod/v4"; import { CredentialItem, vectorStoreCreateCall } from "@/components/networking"; @@ -310,142 +311,129 @@ const VectorStoreForm: React.FC = ({ {selectedProvider === "pg_vector" && ( - -

LiteLLM provides a server to connect to PG Vector. To use this provider:

-
    -
  1. - Deploy the litellm-pgvector server from:{" "} - - https://github.com/BerriAI/litellm-pgvector - -
    2. -
  2. Configure your PostgreSQL database with pgvector extension
    3. -
  3. Start the server and note the API base URL and API key
    4. -
  4. Enter those details in the fields below
    5. -
-
- } - type="info" - showIcon - /> + + + PG Vector Setup Required + +

LiteLLM provides a server to connect to PG Vector. To use this provider:

+
    +
  1. + Deploy the litellm-pgvector server from:{" "} + + https://github.com/BerriAI/litellm-pgvector + +
    2. +
  2. Configure your PostgreSQL database with pgvector extension
    3. +
  3. Start the server and note the API base URL and API key
    4. +
  4. Enter those details in the fields below
    5. +
+ + )} {selectedProvider === "valkey" && ( - -

- LiteLLM searches documents you have already stored in Valkey. It does not create the index or - upload documents for you. Before creating this vector store, make sure: -

-
    -
  1. - Your Valkey server has vector search enabled (the valkey-search module, included in the - valkey-bundle image and in AWS ElastiCache / MemoryDB for Valkey) -
    2. -
  2. - You have already created a search index and loaded your documents and their embeddings into it. - Enter that index name as the Vector Store ID -
    3. -
  3. - You know which embedding model created those stored embeddings. That model must be added to this - proxy under Models so you can pick it below. Using a different model returns wrong results -
    4. -
  4. - You know the field names your documents use for their text and their embedding. If they are not - "text" and "embedding", set them below -
    5. -
-

- When a query comes in, LiteLLM converts it to an embedding with the model below and returns the - closest matching documents from your index. -

-
- } - type="info" - showIcon - /> + + + Valkey Setup Required + +

+ LiteLLM searches documents you have already stored in Valkey. It does not create the index or upload + documents for you. Before creating this vector store, make sure: +

+
    +
  1. + Your Valkey server has vector search enabled (the valkey-search module, included in the + valkey-bundle image and in AWS ElastiCache / MemoryDB for Valkey) +
    2. +
  2. + You have already created a search index and loaded your documents and their embeddings into it. + Enter that index name as the Vector Store ID +
    3. +
  3. + You know which embedding model created those stored embeddings. That model must be added to this + proxy under Models so you can pick it below. Using a different model returns wrong results +
    4. +
  4. + You know the field names your documents use for their text and their embedding. If they are not + "text" and "embedding", set them below +
    5. +
+

+ When a query comes in, LiteLLM converts it to an embedding with the model below and returns the + closest matching documents from your index. +

+ + )} {selectedProvider === "vertex_rag_engine" && ( - -

To use Vertex AI RAG Engine:

-

- Note: Google Cloud has renamed this to "RAG Engine" in its console — the steps below - still apply. -

-
    -
  1. - Set up your Vertex AI RAG Engine corpus following the guide:{" "} - - Vertex AI RAG Engine Overview - -
    2. -
  2. Create a corpus in your Google Cloud project
    3. -
  3. - Note the corpus ID from the Vertex AI console (now labeled "RAG Engine" in Google - Cloud) -
    4. -
  4. Enter the corpus ID in the Vector Store ID field below
    5. -
-
- } - type="info" - showIcon - /> + + + Vertex AI RAG Engine Setup + +

To use Vertex AI RAG Engine:

+

+ Note: Google Cloud has renamed this to "RAG Engine" in its console — the steps below still + apply. +

+
    +
  1. + Set up your Vertex AI RAG Engine corpus following the guide:{" "} + + Vertex AI RAG Engine Overview + +
    2. +
  2. Create a corpus in your Google Cloud project
    3. +
  3. + Note the corpus ID from the Vertex AI console (now labeled "RAG Engine" in Google Cloud) +
    4. +
  4. Enter the corpus ID in the Vector Store ID field below
    5. +
+ + )} {selectedProvider === "vertex_ai/search_api" && ( - -

To use Vertex AI Search (Discovery Engine):

-

- Note: Google Cloud has renamed this to "Agent Search" in its console — the steps below - still apply. -

-
    -
  1. - Enable the Discovery Engine API on your Google Cloud project and create a data store following - the guide:{" "} - - Create a Vertex AI Search data store - -
    2. -
  2. Pick a supported location: global, us, or eu
    3. -
  3. - For most data store types (Cloud Storage, BigQuery, Media): copy the data store ID and enter it - in the Vector Store ID field below. -
    4. -
  4. - For website, healthcare, and connector-based sources (Drive, Gmail, Slack, Jira, etc.): create a - search app on top of the data store, then copy the Engine ID and enter it in - the Engine ID field. The Vector Store ID is still required as the LiteLLM-side name for this - record, but it isn't used in the GCP URL when Engine ID is set. -
    5. -
- - } - type="info" - showIcon - /> + + + Vertex AI Search Setup + +

To use Vertex AI Search (Discovery Engine):

+

+ Note: Google Cloud has renamed this to "Agent Search" in its console — the steps below + still apply. +

+
    +
  1. + Enable the Discovery Engine API on your Google Cloud project and create a data store following the + guide:{" "} + + Create a Vertex AI Search data store + +
    2. +
  2. Pick a supported location: global, us, or eu
    3. +
  3. + For most data store types (Cloud Storage, BigQuery, Media): copy the data store ID and enter it in + the Vector Store ID field below. +
    4. +
  4. + For website, healthcare, and connector-based sources (Drive, Gmail, Slack, Jira, etc.): create a + search app on top of the data store, then copy the Engine ID and enter it in the + Engine ID field. The Vector Store ID is still required as the LiteLLM-side name for this record, + but it isn't used in the GCP URL when Engine ID is set. +
    5. +
+ + )} ; +function SsoEnabledNotice() { + const [isDismissed, setIsDismissed] = useState(false); + if (isDismissed) return null; + + return ( + + + + Single Sign-On (SSO) is enabled. LiteLLM no longer automatically redirects to the SSO login flow upon loading + this page. To re-enable auto-redirect-to-SSO, set{" "} + AUTO_REDIRECT_UI_LOGIN_TO_SSO=true in your + environment configuration. + + + + + + ); +} + function LoginPageContent() { const [isLoading, setIsLoading] = useState(true); const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig(); @@ -170,22 +192,19 @@ function LoginPageContent() {

🚅 LiteLLM

- -

- The Admin UI has been disabled by the administrator. To re-enable it, please update the following - environment variable: -

-

- DISABLE_ADMIN_UI=False -

- - } - type="warning" - showIcon - /> + + + Admin UI Disabled + +

+ The Admin UI has been disabled by the administrator. To re-enable it, please update the following + environment variable: +

+

+ DISABLE_ADMIN_UI=False +

+ + @@ -209,31 +228,32 @@ function LoginPageContent() { {!uiConfig?.hide_default_credentials_hint && ( - -

- By default, Username is admin{" "} - and Password is your set LiteLLM Proxy - MASTER_KEY. -

-

- Need to set UI credentials or SSO?{" "} - - Check the documentation - - . -

- - } - type="info" - icon={} - showIcon - /> + + + Default Credentials + +

+ By default, Username is admin and + Password is your set LiteLLM Proxy + MASTER_KEY. +

+

+ Need to set UI credentials or SSO?{" "} + + Check the documentation + + . +

+ + )} - {error && } + {error && ( + + + {error} + + )}
@@ -326,22 +346,7 @@ function LoginPageContent() {
- {uiConfig?.sso_configured && ( - - Single Sign-On (SSO) is enabled. LiteLLM no longer automatically redirects to the SSO login flow - upon loading this page. To re-enable auto-redirect-to-SSO, set{" "} - AUTO_REDIRECT_UI_LOGIN_TO_SSO=true{" "} - in your environment configuration. - - } - /> - )} + {uiConfig?.sso_configured && } diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx index 1f2441c4870..0a03cc6a591 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx @@ -1,6 +1,7 @@ -import { Alert } from "antd"; +import { CircleAlert, Info } from "lucide-react"; import React from "react"; import { z } from "zod/v4"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Field, FieldLabel, FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; @@ -46,11 +47,10 @@ export function OnboardingFormBody({ variant, userEmail, isPending, claimError,

{variant === "signup" && ( - + + SSO +
SSO is under the Enterprise Tier.
- } - showIcon - /> + + )}
@@ -84,7 +83,12 @@ export function OnboardingFormBody({ variant, userEmail, isPending, claimError, - {claimError && } + {claimError && ( + + + {claimError} + + )}
+ + )} {updateError && ( - + + Could not update settings + {updateError instanceof Error && {updateError.message}} + )} diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 4cb60dc116a..c0c73fbb762 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -5,7 +5,9 @@ import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles"; import { modelCreationScope } from "@/utils/modelPermissions"; import { Switch } from "@/components/ui/switch"; import { Field, FieldLabel } from "@/components/shared/form/field"; -import { Select as AntdSelect, Card, Col, Modal, Row, Tooltip, Typography, Alert } from "antd"; +import { Select as AntdSelect, Card, Col, Modal, Row, Tooltip, Typography } from "antd"; +import { Info } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import type { UploadProps } from "antd/es/upload"; import React, { useEffect, useMemo, useState } from "react"; @@ -168,13 +170,13 @@ const AddModelForm: React.FC = ({ )} {!teamAdminSelectedTeam && ( - + + + Team Selection Required + + As a team admin, you need to select your team first before adding models. + + )} )} diff --git a/ui/litellm-dashboard/src/components/add_pass_through.tsx b/ui/litellm-dashboard/src/components/add_pass_through.tsx index 9f95e7f27ac..09ddac2cf42 100644 --- a/ui/litellm-dashboard/src/components/add_pass_through.tsx +++ b/ui/litellm-dashboard/src/components/add_pass_through.tsx @@ -1,8 +1,9 @@ "use client"; import React, { useState } from "react"; -import { Modal, Alert } from "antd"; -import { CircleHelp, Plug } from "lucide-react"; +import { Modal } from "antd"; +import { CircleHelp, Info, Plug } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { useWatch } from "react-hook-form"; import { z } from "zod/v4"; @@ -156,13 +157,14 @@ const AddPassThroughEndpoint: React.FC = ({ }} >
- + + + What is a Pass-Through Endpoint? + + Route requests from your LiteLLM proxy to any external API. Perfect for custom models, image generation + APIs, or any service you want to proxy through LiteLLM. + + diff --git a/ui/litellm-dashboard/src/components/common_components/PassThroughGuardrailsSection.tsx b/ui/litellm-dashboard/src/components/common_components/PassThroughGuardrailsSection.tsx index f51629b7ee1..a544a8be32b 100644 --- a/ui/litellm-dashboard/src/components/common_components/PassThroughGuardrailsSection.tsx +++ b/ui/litellm-dashboard/src/components/common_components/PassThroughGuardrailsSection.tsx @@ -1,6 +1,6 @@ import React from "react"; -import { Alert } from "antd"; -import { CircleHelp } from "lucide-react"; +import { CircleHelp, Info } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import GuardrailSelector from "../guardrails/GuardrailSelector"; import { TagsInput } from "@/app/(dashboard)/guardrails/_components/content_filter/TagsInput"; @@ -67,21 +67,20 @@ const PassThroughGuardrailsSection: React.FC endpoints.

- - Field-Level Targeting{" "} - - (Learn More) - - - } - description={ + + + + Field-Level Targeting{" "} + + (Learn More) + + +
Optionally specify which fields to check. If left empty, the entire request/response is sent to the @@ -100,11 +99,8 @@ const PassThroughGuardrailsSection: React.FC
- } - type="info" - showIcon - className="mb-4" - /> + + diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx index cbaf00b7131..946bbdd649b 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -74,7 +74,7 @@ describe("UserSearchModal", () => { expect(notice).toHaveTextContent(/users that already exist/i); expect(notice).toHaveTextContent(/ask a proxy admin to create their account first/i); // info, not warning: a warning here would read as an error state on an empty form - expect(notice.className).toMatch(/ant-alert-info/); + expect(notice.className).toMatch(/text-info/); }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index aeead9d82e6..e089b15addc 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -1,5 +1,7 @@ import { useState } from "react"; -import { Modal, Alert } from "antd"; +import { Modal } from "antd"; +import { Info } from "lucide-react"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; import { UserAddOutlined } from "@ant-design/icons"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { useForm } from "react-hook-form"; @@ -202,13 +204,13 @@ const UserSearchModal: React.FC = ({ - + + + + Search selects from users that already exist. To add someone new, ask a proxy admin to create their + account first. + + diff --git a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx index d717c11de86..88554ddba89 100644 --- a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx @@ -1,8 +1,9 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { CheckOutlined, CopyOutlined, SyncOutlined } from "@ant-design/icons"; -import { Alert, Modal, Space } from "antd"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { Modal, Space } from "antd"; import { Button } from "@/components/ui/button"; -import { CircleHelp } from "lucide-react"; +import { CircleHelp, TriangleAlert } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { useWatch } from "react-hook-form"; import { CopyToClipboard } from "react-copy-to-clipboard"; @@ -184,7 +185,10 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat > {regeneratedKey ? (
- + + + Save it now, you will not see it again +
Key Alias diff --git a/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx b/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx index b0a19e7ea05..a436ba55c01 100644 --- a/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx +++ b/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx @@ -1,10 +1,12 @@ -import { Alert, Modal, Typography } from "antd"; +import { Modal, Typography } from "antd"; +import { TriangleAlert } from "lucide-react"; import { useState } from "react"; import { z } from "zod/v4"; import { modelPatchUpdateCall } from "./networking"; import { toast } from "@/lib/toast"; import { FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Button } from "@/components/ui/button"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; @@ -74,12 +76,14 @@ export default function UpdateModelCredentialsModal({ Update this model's API key. Only the new key is sent; the rest of the deployment configuration is left untouched. - + + + + Only the API key is rotated here. Models that authenticate with an Azure AD token, AWS credentials, or a + Vertex service-account JSON aren't supported yet; update those from the model's LiteLLM Params for + now. + + From 74f12bf6efbd931eb28b11015a7998cc95e455fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:09:09 -0700 Subject: [PATCH 30/88] fix(proxy): accept inherited model sentinels in project key limits --- .../key_management_endpoints.py | 10 +++- .../test_key_management_endpoints.py | 50 +++++++++++++++++++ 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 67a836b8c92..ab343cbde2d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1416,6 +1416,11 @@ async def _check_team_key_limits( ) +_INHERITED_MODEL_SENTINELS: Final = frozenset( + {SpecialModelNames.all_team_models.value, SpecialModelNames.all_proxy_models.value} +) + + async def _check_project_key_limits( project_id: str, data: GenerateKeyRequest | UpdateKeyRequest, @@ -1425,7 +1430,8 @@ async def _check_project_key_limits( """ Validate that key's models and budget respect its project's limits. - - Key models must be a subset of project models + - Key models must be a subset of project models, except the all-team-models / all-proxy-models + sentinels, which inherit a parent scope and are narrowed by the project at request time - Key max_budget must be <= project max_budget """ project_obj: Final = await get_project_object( @@ -1443,7 +1449,7 @@ async def _check_project_key_limits( # Validate key models are a subset of project models if data.models and len(project_obj.models) > 0: for m in data.models: - if m not in project_obj.models: + if m not in project_obj.models and m not in _INHERITED_MODEL_SENTINELS: raise HTTPException( status_code=400, detail={ diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdf09a95e4b..0936a2d1439 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( NewUserRequest, LiteLLM_BudgetTable, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -31,9 +32,12 @@ from litellm.proxy._types import ( ResetSpendRequest, UpdateKeyRequest, ) +from litellm.proxy.auth.auth_checks import _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, + _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, _enforce_upperbound_key_params, @@ -16444,3 +16448,49 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( assert await _authorized_models_for_key( access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"] ) == ["attached-model"] + + +async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: + user_api_key_cache = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=_project_cache_key(project_id), + value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), + model_type=LiteLLM_ProjectTableCachedObj, + ) + return user_api_key_cache + + +@pytest.mark.parametrize("request_cls", [GenerateKeyRequest, UpdateKeyRequest]) +@pytest.mark.parametrize("sentinel", ["all-team-models", "all-proxy-models"]) +@pytest.mark.asyncio +async def test_check_project_key_limits_accepts_inherited_model_sentinels(request_cls, sentinel): + """LIT-5823: the sentinels inherit a parent scope, so a project allowlist must not treat them as model names.""" + user_api_key_cache = await _cache_with_project("proj-lit-5823", ["gpt-5.4-nano"]) + + await _check_project_key_limits( + project_id="proj-lit-5823", + data=request_cls(key="sk-lit-5823", models=[sentinel]), + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + ) + + +@pytest.mark.parametrize("request_cls", [GenerateKeyRequest, UpdateKeyRequest]) +@pytest.mark.parametrize( + "key_models", + [["gpt-5.4-mini"], ["all-team-models", "gpt-5.4-mini"], ["gpt-5.4-nano", "all-proxy-models", "gpt-5.4-mini"]], +) +@pytest.mark.asyncio +async def test_check_project_key_limits_still_rejects_real_model_outside_project(request_cls, key_models): + user_api_key_cache = await _cache_with_project("proj-lit-5823", ["gpt-5.4-nano"]) + + with pytest.raises(HTTPException) as exc_info: + await _check_project_key_limits( + project_id="proj-lit-5823", + data=request_cls(key="sk-lit-5823", models=key_models), + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + ) + + assert exc_info.value.status_code == 400 + assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] From fe26f5a541a3b00c72f5b3ecdd4546ec1cf46a6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:11:47 -0700 Subject: [PATCH 31/88] fix(ui): clear the user search spinner when the box is emptied --- .../user_search_modal.test.tsx | 19 +++++++++++++ .../common_components/user_search_modal.tsx | 1 + .../create_key_button.integration.test.tsx | 28 +++++++++++++++++++ .../organisms/create_key_button.tsx | 1 + 4 files changed, 49 insertions(+) diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx index 98b85c738b4..6b71d5ab1d2 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -252,4 +252,23 @@ describe("UserSearchModal out-of-order search results", () => { role: "user", }); }); + + it("stops loading once the box is cleared and the abandoned search answers", async () => { + const user = userEvent.setup(); + render(); + + const input = getEmailSearchInput(); + await user.click(input); + await user.type(input, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + await screen.findByText("Loading..."); + + await user.clear(input); + await screen.findByText("No results"); + + await answerFor("ali", [{ user_id: "u-jones", user_email: "alice.jones@example.com" }]); + + expect(screen.queryByRole("option")).not.toBeInTheDocument(); + expect(screen.getByText("No results")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index 71d38c6db25..41cbd58ba53 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -87,6 +87,7 @@ const UserSearchModal: React.FC = ({ if (!searchText) { setUserOptions([]); + setLoading(false); return; } diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index aa59dbb024f..8072a90fcad 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -835,6 +835,34 @@ describe("CreateKey", () => { expect(screen.queryByTitle("alice.jones@example.com (u-jones)")).not.toBeInTheDocument(); expect(screen.getByTitle("alice.smith@example.com (u-smith)")).toBeInTheDocument(); }); + + it("stops searching once the box is cleared and the abandoned search answers", async () => { + const answers = new Map void>(); + vi.mocked(userFilterUICall).mockImplementation( + (_accessToken, params) => + new Promise((resolve) => { + answers.set(params.get("user_email") ?? "", resolve); + }) as never, + ); + + const user = userEvent.setup(); + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + const search = antdSearchInput(await screen.findByText("Type email to search for users")); + + await user.type(search, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + await screen.findByText("Searching..."); + + await user.clear(search); + await screen.findByText("No users found"); + + await act(async () => { + answers.get("ali")?.([{ user_id: "u-jones", user_email: "alice.jones@example.com" }]); + }); + + expect(screen.queryByTitle("alice.jones@example.com (u-jones)")).not.toBeInTheDocument(); + expect(screen.getByText("No users found")).toBeInTheDocument(); + }); }); describe("created key display", () => { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 73438c244c1..62ad62e0f83 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -563,6 +563,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp if (!searchText) { setUserOptions([]); + setUserSearchLoading(false); return; } From 77716eeaed12e83bbfbfb88282af8ea92eb81c2b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:13:00 -0700 Subject: [PATCH 32/88] fix(model_prices): set prompt_cache_min_tokens=4096 for Gemini 3.5/3.6/3.7 Flash and 3.1 Pro Preview --- ...odel_prices_and_context_window_backup.json | 15 +++++++ model_prices_and_context_window.json | 15 +++++++ tests/test_litellm/test_utils.py | 40 +++++++++++++++++++ 3 files changed, 70 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f9027313b..8c07ca35443 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19403,6 +19403,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19460,6 +19461,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19614,6 +19616,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19665,6 +19668,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -19719,6 +19723,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -19773,6 +19778,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19830,6 +19836,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -21337,6 +21344,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -21391,6 +21399,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21448,6 +21457,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21538,6 +21548,7 @@ "tpm": 800000 }, "gemini/gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -21595,6 +21606,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -21733,6 +21745,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -21785,6 +21798,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21840,6 +21854,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 07f9027313b..8c07ca35443 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19403,6 +19403,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19460,6 +19461,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19614,6 +19616,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19665,6 +19668,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -19719,6 +19723,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -19773,6 +19778,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19830,6 +19836,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -21337,6 +21344,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -21391,6 +21399,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21448,6 +21457,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21538,6 +21548,7 @@ "tpm": 800000 }, "gemini/gemini-3.1-pro-preview": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -21595,6 +21606,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-pro-preview-customtools": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -21733,6 +21745,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -21785,6 +21798,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.6-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, @@ -21840,6 +21854,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.7-flash": { + "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "input_cost_per_token": 7.5e-07, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index efccdc4a986..4e75859c6df 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2,6 +2,7 @@ import json import logging import os import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -4384,6 +4385,45 @@ def test_get_prompt_cache_min_tokens_differs_per_platform_for_same_model(local_m ) +GEMINI_4096_CACHE_MIN_MODELS: Final = tuple( + prefix + base + for base in ( + "gemini-3.5-flash", + "gemini-3.6-flash", + "gemini-3.7-flash", + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview-customtools", + ) + for prefix in ("", "gemini/", "vertex_ai/") +) + + +def test_gemini_3_flash_and_31_pro_preview_resolve_4096_cache_minimum(local_model_cost_map: None) -> None: + """Regression for the cost map missing prompt_cache_min_tokens on these models: Google rejects + explicit caching below 4,096 tokens for them (https://ai.google.dev/gemini-api/docs/caching), so + the 1024 default sent cachedContents creates Vertex answered with a hard 400.""" + wrong: Final = { + model: get_prompt_cache_min_tokens(model=model) + for model in GEMINI_4096_CACHE_MIN_MODELS + if get_prompt_cache_min_tokens(model=model) != 4096 + } + assert not wrong, f"prompt_cache_min_tokens must be 4096: {wrong}" + + +def test_gemini_4096_cache_minimum_present_in_root_cost_map() -> None: + """The root map ships to the CDN independently of the bundled backup, so both must carry the + minimum or proxies reading one of them regress to the 1024 default.""" + root_map_path: Final = os.path.join(os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json") + with open(root_map_path) as f: + root_map: Final = json.load(f) + wrong: Final = { + model: root_map[model].get("prompt_cache_min_tokens") + for model in GEMINI_4096_CACHE_MIN_MODELS + if root_map[model].get("prompt_cache_min_tokens") != 4096 + } + assert not wrong, f"prompt_cache_min_tokens must be 4096: {wrong}" + + def test_get_prompt_cache_min_tokens_unmapped_model_falls_back_to_default(local_model_cost_map: None) -> None: """get_model_info raises for a model it has no entry for. The resolver must swallow that and fall back to the default, otherwise the raise reaches callers that would read it as From f1e143a87c34b01104e2ec0dc15be554355d45b3 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 14:18:08 -0700 Subject: [PATCH 33/88] chore(ui): upgrade the dashboard to React 19 (#37411) * chore(ui): upgrade the dashboard to React 19 Bumps react and react-dom from 18.3.1 to 19.2.8 with matching @types. Next 16 already required a React 19 peer, so this aligns the dashboard with what the framework expects and unblocks Base UI and shadcn work that assumes the React 19 ref model. React 19 passes ref through as a regular prop, so the setup file's forwardRef tripwire and the ref-forwarding test's forwardRef case no longer describe real behavior; both now assert the React 19 contract instead. useRef(null) now yields RefObject, which is the one prop type MessageList had to widen. * test(ui): wait for a Base UI select popup to open before clicking an option The option lands in the DOM one render before the popup finishes entering, while its positioner still carries pointer-events: none, so clicking it throws. Waiting on the option's text alone was a race that React 19's flush timing loses, which is why four ToolPolicies cases went red on the bump. chooseSelectOption in test-utils opens the trigger, finds the option by role, waits for it to stop being pointer-blocked, then clicks. It also replaces the last-match-by-text hack, which only worked because the popup happens to portal after the table. --- ui/litellm-dashboard/package-lock.json | 157 ++++++++---------- ui/litellm-dashboard/package.json | 8 +- .../_components/mcp_server_edit.test.tsx | 18 ++ .../_components/mcp_server_edit.tsx | 5 +- .../conversation_panel/MessageList.tsx | 2 +- .../UsageViewSelect/UsageViewSelect.test.tsx | 5 +- .../ToolPolicies/ToolPoliciesPanel.test.tsx | 31 ++-- .../src/components/team/TeamInfo.test.tsx | 8 +- .../src/components/ui/ref-forwarding.test.tsx | 16 +- ui/litellm-dashboard/tests/setupTests.ts | 26 +-- ui/litellm-dashboard/tests/test-utils.tsx | 29 +++- 11 files changed, 146 insertions(+), 159 deletions(-) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 186d38234d5..7069450a75c 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -32,9 +32,9 @@ "openapi-fetch": "^0.17.0", "openapi-react-query": "^0.5.4", "papaparse": "5.5.3", - "react": "18.3.1", + "react": "19.2.8", "react-copy-to-clipboard": "5.1.1", - "react-dom": "18.3.1", + "react-dom": "19.2.8", "react-hook-form": "7.82.0", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", @@ -55,9 +55,9 @@ "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", "@types/node": "20.19.37", - "@types/react": "18.2.48", + "@types/react": "19.2.18", "@types/react-copy-to-clipboard": "5.0.7", - "@types/react-dom": "18.3.7", + "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", "@vitest/coverage-v8": "3.2.6", "@vitest/ui": "3.2.6", @@ -200,9 +200,9 @@ } }, "node_modules/@ant-design/icons-svg": { - "version": "4.4.2", - "resolved": "https://registry.npmjs.org/@ant-design/icons-svg/-/icons-svg-4.4.2.tgz", - "integrity": "sha512-vHbT+zJEVzllwP+CM+ul7reTEfBR0vgxFe7+lREAsAA7YGsYpboiq2sQNeQeRvh09GfQgs/GyFEvZpJ9cLXpXA==", + "version": "4.5.0", + "resolved": "https://registry.npmjs.org/@ant-design/icons-svg/-/icons-svg-4.5.0.tgz", + "integrity": "sha512-1BTUFyKPTBZ53MuTP8s0k5SFEXL7o3VHEOwLgzaoWKwnBeqIcqUtVshc4SKzhI6uACfqhJqBwBUE9FsWR3uULA==", "license": "MIT" }, "node_modules/@ant-design/react-slick": { @@ -606,19 +606,6 @@ } } }, - "node_modules/@base-ui/react/node_modules/@floating-ui/react-dom": { - "version": "2.1.8", - "resolved": "https://registry.npmjs.org/@floating-ui/react-dom/-/react-dom-2.1.8.tgz", - "integrity": "sha512-cC52bHwM/n/CxS87FH0yWdngEZrjdtLW/qVruo68qg+prK7ZQ4YGdut2GyDVpoGeAYe/h899rVeOVm6Oi40k2A==", - "license": "MIT", - "dependencies": { - "@floating-ui/dom": "^1.7.6" - }, - "peerDependencies": { - "react": ">=16.8.0", - "react-dom": ">=16.8.0" - } - }, "node_modules/@base-ui/utils": { "version": "0.3.1", "resolved": "https://registry.npmjs.org/@base-ui/utils/-/utils-0.3.1.tgz", @@ -1441,28 +1428,41 @@ } }, "node_modules/@floating-ui/core": { - "version": "1.7.5", - "resolved": "https://registry.npmjs.org/@floating-ui/core/-/core-1.7.5.tgz", - "integrity": "sha512-1Ih4WTWyw0+lKyFMcBHGbb5U5FtuHJuujoyyr5zTaWS5EYMeT6Jb2AuDeftsCsEuchO+mM2ij5+q9crhydzLhQ==", + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@floating-ui/core/-/core-1.8.0.tgz", + "integrity": "sha512-0CIZ5itps/8x7BG8dEIhs53BvCUH2PCoogtakwRTut+Arm58sJooJ0AuZhLw2HJYIR5cMLNPBSS728sPho2khQ==", "license": "MIT", "dependencies": { - "@floating-ui/utils": "^0.2.11" + "@floating-ui/utils": "^0.2.12" } }, "node_modules/@floating-ui/dom": { - "version": "1.7.6", - "resolved": "https://registry.npmjs.org/@floating-ui/dom/-/dom-1.7.6.tgz", - "integrity": "sha512-9gZSAI5XM36880PPMm//9dfiEngYoC6Am2izES1FF406YFsjvyBMmeJ2g4SAju3xWwtuynNRFL2s9hgxpLI5SQ==", + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@floating-ui/dom/-/dom-1.8.0.tgz", + "integrity": "sha512-yXSrzeHZBTZadLOlfyhCkJHNeLJnHRnRInwdZ40L7ZiaAtrBwoYlsDrX3v5zB1Utk7CLfzcOVnVVWoXEky7Ceg==", "license": "MIT", "dependencies": { - "@floating-ui/core": "^1.7.5", - "@floating-ui/utils": "^0.2.11" + "@floating-ui/core": "^1.8.0", + "@floating-ui/utils": "^0.2.12" + } + }, + "node_modules/@floating-ui/react-dom": { + "version": "2.1.9", + "resolved": "https://registry.npmjs.org/@floating-ui/react-dom/-/react-dom-2.1.9.tgz", + "integrity": "sha512-JDjEFGCpImxDCA7JJKviA0M9+RtmJdj0m/NVU5IMgBK+AmZouAQQ7/+2GLH0GXXY0YMw9oXPB8hKdbPYg5QLYg==", + "license": "MIT", + "dependencies": { + "@floating-ui/dom": "^1.8.0" + }, + "peerDependencies": { + "react": ">=16.8.0", + "react-dom": ">=16.8.0" } }, "node_modules/@floating-ui/utils": { - "version": "0.2.11", - "resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.11.tgz", - "integrity": "sha512-RiB/yIh78pcIxl6lLMG0CgBXAZ2Y0eVHqMPYugu+9U0AeT6YBeiJpf7lbdJNIugFP5SIjwNRgo4DhR1Qxi26Gg==", + "version": "0.2.12", + "resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.12.tgz", + "integrity": "sha512-HpCo8tmWzLVad5s2d19EhAz5zqrrQ6s69qd6moPMQvkOuSwDT1YgRfWSVuc4ennqrgv3OHppiOGMQ7oC13yIww==", "license": "MIT" }, "node_modules/@headlessui/tailwindcss": { @@ -2627,9 +2627,9 @@ "license": "MIT" }, "node_modules/@rc-component/async-validator": { - "version": "5.1.0", - "resolved": "https://registry.npmjs.org/@rc-component/async-validator/-/async-validator-5.1.0.tgz", - "integrity": "sha512-n4HcR5siNUXRX23nDizbZBQPO0ZM/5oTtmKZ6/eqL0L2bo747cklFdZGRN2f+c9qWGICwDzrhW0H7tE9PptdcA==", + "version": "5.1.2", + "resolved": "https://registry.npmjs.org/@rc-component/async-validator/-/async-validator-5.1.2.tgz", + "integrity": "sha512-WYbrZSjzznU1ekD0qFq2qRxt309VoS61MTG5npnFQlKYcoy9IzU8T+ZCIhq5bGAXRbXysABFWTspicMfmWFwow==", "license": "MIT", "dependencies": { "@babel/runtime": "^7.24.4" @@ -2669,9 +2669,9 @@ } }, "node_modules/@rc-component/mini-decimal": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@rc-component/mini-decimal/-/mini-decimal-1.1.3.tgz", - "integrity": "sha512-bk/FJ09fLf+NLODMAFll6CfYrHPBioTedhW6lxDBuuWucJEqFUd4l/D/5JgIi3dina6sYahB8iuPAZTNz2pMxw==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rc-component/mini-decimal/-/mini-decimal-1.1.4.tgz", + "integrity": "sha512-xiuXcaCwyOWpD8a8scdExFl+bntNphAW8XeenL1ig2en0AAZY0Pcp4pC0dI22qJ+NvxKn9RoNIoRdqYU3BLH4w==", "license": "MIT", "dependencies": { "@babel/runtime": "^7.18.0" @@ -2717,9 +2717,9 @@ } }, "node_modules/@rc-component/qrcode": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/@rc-component/qrcode/-/qrcode-1.1.1.tgz", - "integrity": "sha512-LfLGNymzKdUPjXUbRP+xOhIWY4jQ+YMj5MmWAcgcAq1Ij8XP7tRmAXqyuv96XvLUBE/5cA8hLFl9eO1JQMujrA==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rc-component/qrcode/-/qrcode-1.1.3.tgz", + "integrity": "sha512-aGv6alnn4HbDEsURzKP+jv13rbi1VxmAYfBNZr5GKF1iohMNWy5tAVoJ1E3cOvzMB1kbUPvCXchM6zSFlRGPhA==", "license": "MIT", "dependencies": { "@babel/runtime": "^7.24.7" @@ -3661,12 +3661,12 @@ } }, "node_modules/@tanstack/react-store": { - "version": "0.11.0", - "resolved": "https://registry.npmjs.org/@tanstack/react-store/-/react-store-0.11.0.tgz", - "integrity": "sha512-tX4YXh3PDkmpvGQWkWqKpzs/MSqbtuwY9dWdWhtV9Q50PmO+jOkUKIWIX4G85dwt7lxdHLXsiaEKPdKmC8F41w==", + "version": "0.11.1", + "resolved": "https://registry.npmjs.org/@tanstack/react-store/-/react-store-0.11.1.tgz", + "integrity": "sha512-HaIGKI3YLmjBYIvy5DFDY23oNaYZIsTZfngey07Uh5iLVJgM3bIGCnZeOFOqzjFld9JHWcaHJnasD/bKoGKwJQ==", "license": "MIT", "dependencies": { - "@tanstack/store": "0.11.0", + "@tanstack/store": "0.11.1", "use-sync-external-store": "^1.6.0" }, "funding": { @@ -3699,9 +3699,9 @@ } }, "node_modules/@tanstack/store": { - "version": "0.11.0", - "resolved": "https://registry.npmjs.org/@tanstack/store/-/store-0.11.0.tgz", - "integrity": "sha512-WlzzCt3xi0G6pCAJu1U+2jiECwabETDpQDi3hfkFZvJii9AuZqEKbOiVarX1/bWhTNjU486yQtJCCasi/0q+Cw==", + "version": "0.11.1", + "resolved": "https://registry.npmjs.org/@tanstack/store/-/store-0.11.1.tgz", + "integrity": "sha512-mzTOBhypOuDJAy/D8n2MfUZ1HFkXnmSETviRyhqEC8LUE7/IZQExOTxMANj3KjTofYTkFNpBY67qaVrT41YccA==", "license": "MIT", "funding": { "type": "github", @@ -3999,21 +3999,13 @@ "@types/node": "*" } }, - "node_modules/@types/prop-types": { - "version": "15.7.15", - "resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.15.tgz", - "integrity": "sha512-F6bEyamV9jKGAFBEmlQnesRPGOQqS2+Uwi0Em15xenOxHaf2hv6L8YCVn3rPdPJOiJfPiCnLIRyvwVaqMY3MIw==", - "license": "MIT" - }, "node_modules/@types/react": { - "version": "18.2.48", - "resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz", - "integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==", + "version": "19.2.18", + "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.18.tgz", + "integrity": "sha512-AnzbBERsrLKtk2XSfTbYRLjQPdy116Sty4q+T+Bp3IC4l6jNBvreVPAHmpq9qhXQM7CXZPjLVmGMw9sy+hxQ3w==", "license": "MIT", "dependencies": { - "@types/prop-types": "*", - "@types/scheduler": "*", - "csstype": "^3.0.2" + "csstype": "^3.2.2" } }, "node_modules/@types/react-copy-to-clipboard": { @@ -4027,13 +4019,13 @@ } }, "node_modules/@types/react-dom": { - "version": "18.3.7", - "resolved": "https://registry.npmjs.org/@types/react-dom/-/react-dom-18.3.7.tgz", - "integrity": "sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==", + "version": "19.2.4", + "resolved": "https://registry.npmjs.org/@types/react-dom/-/react-dom-19.2.4.tgz", + "integrity": "sha512-Bsc+QHgp+P/F02XDzNCY9jnZNCUuLki36KT7VKrTXXLdHf+vHMNZnW1rVu5DNW/rCK+fya3DATySbLM4yhtKUw==", "dev": true, "license": "MIT", "peerDependencies": { - "@types/react": "^18.0.0" + "@types/react": "^19.2.0" } }, "node_modules/@types/react-syntax-highlighter": { @@ -4046,12 +4038,6 @@ "@types/react": "*" } }, - "node_modules/@types/scheduler": { - "version": "0.26.0", - "resolved": "https://registry.npmjs.org/@types/scheduler/-/scheduler-0.26.0.tgz", - "integrity": "sha512-WFHp9YUJQ6CKshqoC37iOlHnQSmxNc795UhB26CyBBttrN9svdIrUjl/NjnNmfcwtncN0h/0PPAFWv9ovP8mLA==", - "license": "MIT" - }, "node_modules/@types/unist": { "version": "3.0.3", "resolved": "https://registry.npmjs.org/@types/unist/-/unist-3.0.3.tgz", @@ -11565,13 +11551,10 @@ } }, "node_modules/react": { - "version": "18.3.1", - "resolved": "https://registry.npmjs.org/react/-/react-18.3.1.tgz", - "integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==", + "version": "19.2.8", + "resolved": "https://registry.npmjs.org/react/-/react-19.2.8.tgz", + "integrity": "sha512-PWaYA1L/q9u2u7xYQi+Y3L3Yfnie7XyLeaJICV1MGD6LprsBxcAqGjYyr0eY3p+QdsA+x/Irkt4Qif8D63+Sbw==", "license": "MIT", - "dependencies": { - "loose-envify": "^1.1.0" - }, "engines": { "node": ">=0.10.0" } @@ -11590,16 +11573,15 @@ } }, "node_modules/react-dom": { - "version": "18.3.1", - "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz", - "integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==", + "version": "19.2.8", + "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.8.tgz", + "integrity": "sha512-rVprimfGBG3DR+Tq0IQG2DT5PxKth1WIGDmj5yPmlzr4YBe7uyE+Du4oVqTDXZSHGGGXRtTJEGSSePyQCMBglQ==", "license": "MIT", "dependencies": { - "loose-envify": "^1.1.0", - "scheduler": "^0.23.2" + "scheduler": "^0.27.0" }, "peerDependencies": { - "react": "^18.3.1" + "react": "^19.2.8" } }, "node_modules/react-hook-form": { @@ -12204,13 +12186,10 @@ } }, "node_modules/scheduler": { - "version": "0.23.2", - "resolved": "https://registry.npmjs.org/scheduler/-/scheduler-0.23.2.tgz", - "integrity": "sha512-UOShsPwz7NrMUqhR6t0hWjFduvOzbtv7toDH1/hIrfRNIDBnnBWd0CwJTGvTpngVlmwGCdP9/Zl/tVrDqcuYzQ==", - "license": "MIT", - "dependencies": { - "loose-envify": "^1.1.0" - } + "version": "0.27.0", + "resolved": "https://registry.npmjs.org/scheduler/-/scheduler-0.27.0.tgz", + "integrity": "sha512-eNv+WrVbKu1f3vbYJT/xtiF5syA5HPIMtf9IgY/nKg0sWqzAUEvqY/xm7OcZc/qafLx/iO9FgOmeSAp4v5ti/Q==", + "license": "MIT" }, "node_modules/scroll-into-view-if-needed": { "version": "3.1.0", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ac13a12620d..962cabcba7b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -48,9 +48,9 @@ "openapi-fetch": "^0.17.0", "openapi-react-query": "^0.5.4", "papaparse": "5.5.3", - "react": "18.3.1", + "react": "19.2.8", "react-copy-to-clipboard": "5.1.1", - "react-dom": "18.3.1", + "react-dom": "19.2.8", "react-hook-form": "7.82.0", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", @@ -71,9 +71,9 @@ "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", "@types/node": "20.19.37", - "@types/react": "18.2.48", + "@types/react": "19.2.18", "@types/react-copy-to-clipboard": "5.0.7", - "@types/react-dom": "18.3.7", + "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", "@vitest/coverage-v8": "3.2.6", "@vitest/ui": "3.2.6", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index 14e26fa5393..96685bb8359 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -1939,6 +1939,24 @@ describe("MCPServerEdit OAuth flow prefill display", () => { expect(screen.queryByText("This server has no OAuth flow set")).not.toBeInTheDocument(); }); + it("never flashes the warning while mounting a server that already has a flow", async () => { + const flashes: Node[] = []; + const observer = new MutationObserver((records) => { + for (const record of records) { + for (const node of record.addedNodes) { + if (node.textContent?.includes("This server has no OAuth flow set")) flashes.push(node); + } + } + }); + observer.observe(document.body, { childList: true, subtree: true }); + + renderEdit({ oauth2_flow: "client_credentials" }); + await screen.findByText("Machine-to-Machine (M2M)"); + observer.disconnect(); + + expect(flashes).toHaveLength(0); + }); + it("does not warn for a delegate (PKCE passthrough) server even with no flow set", () => { renderEdit({ oauth2_flow: null, delegate_auth_to_upstream: true }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 94b66f1bdfa..4b092aa3f80 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -173,7 +173,10 @@ const MCPServerEdit: React.FC = ({ const isTokenExchangeAuthType = authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE; const isIdJagAuthType = authType === AUTH_TYPE.OAUTH2_ID_JAG; const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4; - const oauthFlowTypeValue = mountedValues.oauth_flow_type as string | undefined; + // Same fallback as the delegate switch below: the value is undefined until the field mounts, so + // reading it alone flashes the "no OAuth flow set" warning at a server that already has one. + const oauthFlowTypeValue = + (mountedValues.oauth_flow_type as string | undefined) ?? oauth2FlowToFormValue(mcpServer.oauth2_flow); const isM2MFlow = isOAuthAuthType && oauthFlowTypeValue === OAUTH_FLOW.M2M; // Watch reflects a live toggle when the delegate switch is mounted; fall back to // the stored value otherwise (useWatch returns undefined for an unmounted field, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageList.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageList.tsx index 347553630ae..48d6697a052 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageList.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/MessageList.tsx @@ -8,7 +8,7 @@ interface MessageListProps { messages: Message[]; isLoading: boolean; hasVariables: boolean; - messagesEndRef: React.RefObject; + messagesEndRef: React.RefObject; } const MessageList: React.FC = ({ messages, isLoading, hasVariables, messagesEndRef }) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx index 3d835c98f52..8848ef9f49e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx @@ -1,6 +1,7 @@ import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { chooseSelectOption } from "@/../tests/test-utils"; import { UsageViewSelect } from "./UsageViewSelect"; const openMenu = async (user: ReturnType) => { @@ -35,9 +36,7 @@ describe("UsageViewSelect", () => { const user = userEvent.setup(); render(); - await openMenu(user); - const matches = screen.getAllByText("Team Usage"); - await user.click(matches[matches.length - 1]); + await chooseSelectOption(user, screen.getByRole("combobox"), /^Team Usage/); expect(mockOnChange).toHaveBeenCalled(); expect(mockOnChange.mock.calls[0][0]).toBe("team"); diff --git a/ui/litellm-dashboard/src/components/ToolPolicies/ToolPoliciesPanel.test.tsx b/ui/litellm-dashboard/src/components/ToolPolicies/ToolPoliciesPanel.test.tsx index aaa70a5e357..a2d3539a58b 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies/ToolPoliciesPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies/ToolPoliciesPanel.test.tsx @@ -4,7 +4,7 @@ import { act, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { focusManager, QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import type { ToolRow } from "@/components/networking"; import { ToolPoliciesPanel } from "./ToolPoliciesPanel"; import { toast } from "@/lib/toast"; @@ -84,17 +84,6 @@ const policyValue = (toolId: string, kind: "input" | "output"): string => { const isSaving = (toolId: string, kind: "input" | "output"): boolean => policySelect(toolId, kind).hasAttribute("disabled"); -const chooseOption = async (user: ReturnType, trigger: HTMLElement, label: string) => { - await user.click(trigger); - // The label also renders in the trigger once selected, so take the last match: - // the popup is portalled after the table in document order. - const option = await waitFor(() => { - const matches = screen.getAllByText(label); - return matches[matches.length - 1]; - }); - await user.click(option); -}; - const renderPanel = (onSelectTool = vi.fn()) => renderWithProviders(); @@ -202,7 +191,7 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); expect(updateToolPolicy).toHaveBeenCalledWith("sk-token", "get_weather", { input_policy: "trusted" }); await waitFor(() => expect(policyValue("tool-1", "input")).toBe("trusted")); @@ -214,7 +203,7 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "output"), "trusted"); + await chooseSelectOption(user, policySelect("tool-1", "output"), "trusted"); expect(updateToolPolicy).toHaveBeenCalledWith("sk-token", "get_weather", { output_policy: "trusted" }); }); @@ -225,8 +214,8 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); - await chooseOption(user, policySelect("tool-2", "input"), "blocked"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-2", "input"), "blocked"); expect(isSaving("tool-2", "input")).toBe(true); expect(isSaving("tool-1", "input")).toBe(true); @@ -241,8 +230,8 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); - await chooseOption(user, policySelect("tool-2", "input"), "blocked"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-2", "input"), "blocked"); await act(async () => { finishFirst(); }); @@ -262,7 +251,7 @@ describe("ToolPoliciesPanel inline policy editing", () => { () => new Promise((resolve) => (landStaleRefresh = () => resolve(TOOLS))), ); await user.click(screen.getByTestId("datatable-refresh")); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); await waitFor(() => expect(policyValue("tool-1", "input")).toBe("trusted")); await act(async () => { @@ -281,7 +270,7 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); await waitFor(() => expect(fromBackend).toHaveBeenCalledWith("Failed to update input policy: nope")); expect(policyValue("tool-1", "input")).toBe("untrusted"); @@ -293,7 +282,7 @@ describe("ToolPoliciesPanel inline policy editing", () => { renderPanel(); await waitForRows(); - await chooseOption(user, policySelect("tool-1", "input"), "trusted"); + await chooseSelectOption(user, policySelect("tool-1", "input"), "trusted"); await waitFor(() => expect(isSaving("tool-1", "input")).toBe(true)); expect(isSaving("tool-1", "output")).toBe(false); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 77f35762528..a7e8e6788dd 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -3,7 +3,7 @@ import * as networking from "@/components/networking"; import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import TeamInfoView from "./TeamInfo"; const authState = vi.hoisted(() => ({ userRole: "Admin" })); @@ -993,8 +993,7 @@ describe("TeamInfoView", () => { const user = userEvent.setup({ delay: null }); const resetBudgetSelect = await openSettingsEditorForTeam(user, { budget_duration: "30d" }); - await user.click(resetBudgetSelect); - await user.click(await screen.findByText("Never resets")); + await chooseSelectOption(user, resetBudgetSelect, "Never resets"); await waitFor(() => { expect(resetBudgetSelect).toHaveTextContent("Never resets"); @@ -1028,8 +1027,7 @@ describe("TeamInfoView", () => { const user = userEvent.setup({ delay: null }); const resetBudgetSelect = await openSettingsEditorForTeam(user, { budget_duration: null }); - await user.click(resetBudgetSelect); - await user.click(await screen.findByText("weekly")); + await chooseSelectOption(user, resetBudgetSelect, "weekly"); await user.click(screen.getByRole("button", { name: /save changes/i })); diff --git a/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx b/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx index 48b1e8be226..66f7bd27138 100644 --- a/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx +++ b/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx @@ -138,15 +138,13 @@ describe("ui primitives forward refs to their DOM node", () => { }); }); -describe("setupTests ref tripwire", () => { - it("records a violation when a ref is passed to a plain function component", () => { - const Plain = (props: React.ComponentPropsWithoutRef<"span">) => ; +describe("plain function components", () => { + it("receives a ref as a prop instead of dropping it", () => { + const Plain = (props: React.ComponentPropsWithoutRef<"span"> & { ref?: React.Ref }) => ( + + ); const ref = React.createRef(); - render(React.createElement(Plain as never, { ref })); - const consume = (globalThis as { __consumePendingRefWarnings?: () => string[] }).__consumePendingRefWarnings; - expect(consume).toBeDefined(); - const violations = consume!(); - expect(violations).toHaveLength(1); - expect(violations[0]).toContain("Function components cannot be given refs"); + render(ok); + expect(ref.current).toBeInstanceOf(HTMLSpanElement); }); }); diff --git a/ui/litellm-dashboard/tests/setupTests.ts b/ui/litellm-dashboard/tests/setupTests.ts index 3eceef211a5..d007206e681 100644 --- a/ui/litellm-dashboard/tests/setupTests.ts +++ b/ui/litellm-dashboard/tests/setupTests.ts @@ -113,31 +113,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ }), })); -const pendingRefWarnings: string[] = []; -const consumePendingRefWarnings = (): string[] => pendingRefWarnings.splice(0, pendingRefWarnings.length); -(globalThis as { __consumePendingRefWarnings?: () => string[] }).__consumePendingRefWarnings = - consumePendingRefWarnings; - -const originalConsoleError = console.error.bind(console); -vi.spyOn(console, "error").mockImplementation((...args: unknown[]) => { - originalConsoleError(...args); - if (typeof args[0] === "string" && args[0].includes("Function components cannot be given refs")) { - pendingRefWarnings.push(args.map(String).join(" ")); - } -}); - -afterEach(() => { - cleanup(); - const refWarnings = consumePendingRefWarnings(); - if (refWarnings.length > 0) { - throw new Error( - "A ref was passed to a plain function component and silently dropped under React 18, which breaks " + - "ref-based composition (Base UI render triggers, tooltips, focus). Wrap the component in React.forwardRef. " + - "This tripwire lives in tests/setupTests.ts and can be removed after the React 19 upgrade.\n\n" + - refWarnings.join("\n\n"), - ); - } -}); +afterEach(cleanup); // Make toLocaleString deterministic in tests; individual tests can override // This returns ISO-like strings to keep assertions stable. diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index 2573fbdfeac..66966201a9c 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -1,7 +1,9 @@ import React, { PropsWithChildren } from "react"; -import { render, RenderOptions } from "@testing-library/react"; +import { render, RenderOptions, screen, waitFor } from "@testing-library/react"; +import type userEvent from "@testing-library/user-event"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { NuqsTestingAdapter, OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import { expect } from "vitest"; // Create a client for testing export const testQueryClient = new QueryClient({ @@ -35,4 +37,29 @@ export const renderWithProviders = (ui: React.ReactElement, options?: RenderOpti return render(ui, { wrapper: Providers, ...renderOptions }); }; +const pointerBlocked = (element: HTMLElement): boolean => { + for (let node: HTMLElement | null = element; node !== null; node = node.parentElement) { + if (node.style.pointerEvents === "none") return true; + } + return false; +}; + +/** + * Opens a Base UI Select and picks an option by its accessible name. + * + * The option is in the DOM one render before the popup finishes entering, and until then its + * positioner still carries `pointer-events: none`, which user-event refuses to click. Waiting on + * the option text alone is a race that React 19's flush timing loses. + */ +export const chooseSelectOption = async ( + user: ReturnType, + trigger: HTMLElement, + optionName: string | RegExp, +) => { + await user.click(trigger); + const option = await screen.findByRole("option", { name: optionName }); + await waitFor(() => expect(pointerBlocked(option)).toBe(false)); + await user.click(option); +}; + export * from "@testing-library/react"; From afbfc3f8fa8391a232b48c66e1d715bf9477888c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 19 Aug 2026 14:19:24 -0700 Subject: [PATCH 34/88] fix(complexity-router): gate the reasoning override on a non-SIMPLE score (#37500) Two or more reasoning keyword matches promoted a request straight to the REASONING tier no matter what the weighted score said, so "hi, step by step, pros and cons" scored 0.100 and still bought the most expensive tier. Require the score to clear the simple_medium boundary before the override applies. Promotion from MEDIUM or COMPLEX is unchanged; only prompts the scorer already placed in the cheapest band stay there. --- .../complexity_router/complexity_router.py | 7 ++++--- .../router_strategy/test_complexity_router.py | 18 ++++++++++++++++++ .../add_model/ClassificationMethodConfig.tsx | 2 +- .../RoutingDecisionCard.test.tsx | 12 +++++++++--- .../LogDetailsDrawer/RoutingDecisionCard.tsx | 2 +- 5 files changed, 33 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index d16063b9bd4..0573f8acf18 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1020,13 +1020,14 @@ class ComplexityRouter(CustomLogger): weights: Final = self.config.dimension_weights weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) - # Check for reasoning override (2+ reasoning markers) + boundaries: Final = self._effective_tier_boundaries() + scored_above_simple: Final = weighted_score >= boundaries["simple_medium"] + # Reuse match count from _score_keyword_match to avoid scanning twice - if reasoning_match_count >= 2: + if reasoning_match_count >= 2 and scored_above_simple: return ComplexityTier.REASONING, weighted_score, tuple(signals), "reasoning_override" # Map score to tier - boundaries: Final = self._effective_tier_boundaries() if weighted_score < boundaries["simple_medium"]: tier = ComplexityTier.SIMPLE elif weighted_score < boundaries["medium_complex"]: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index e1e8d9553b3..a148da8b675 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -267,6 +267,24 @@ class TestReasoningMarkerScoring: # 2+ reasoning markers should force REASONING tier assert tier == ComplexityTier.REASONING + def test_reasoning_override_does_not_rescue_a_simple_score(self, complexity_router): + """Reasoning markers on an otherwise trivial prompt must not reach REASONING.""" + prompt = "hi, step by step, pros and cons" + tier, score, signals = complexity_router.classify(prompt) + assert score < complexity_router.config.tier_boundaries["simple_medium"] + assert any("step by step" in s and "pros and cons" in s for s in signals) + assert tier == ComplexityTier.SIMPLE + + def test_reasoning_override_applies_at_the_simple_medium_boundary(self, complexity_router): + """A score sitting exactly on simple_medium is not SIMPLE, so the override still promotes it.""" + prompt = ( + "Give me the pros and cons, step by step, of moving our checkout service " + "to an event-driven architecture." + ) + tier, score, signals = complexity_router.classify(prompt) + assert score == complexity_router.config.tier_boundaries["simple_medium"] + assert tier == ComplexityTier.REASONING + def test_system_prompt_reasoning_not_counted(self, complexity_router): """Reasoning markers in system prompt should not count for override.""" user_prompt = "What is 2+2?" diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 8d42b34a14b..da0c50c7bd0 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -93,7 +93,7 @@ const HowClassificationWorks: React.FC<{ value: ComplexityRouterConfigValue }> =
  • {effectiveTierLabel("REASONING", value.tier_labels)}: Score > {ranges.complexReasoning}{" "} - (or 2+ reasoning markers) + (or 2+ reasoning markers with a score of at least {ranges.simpleMedium})
    • )} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index a2afe0d7a0d..ed99e414706 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -59,7 +59,9 @@ describe("RoutingDecisionCard", () => { }} />, ); - expect(screen.getByText("Heuristic, REASONING override (2 or more reasoning markers)")).toBeInTheDocument(); + expect( + screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score above the lowest tier)"), + ).toBeInTheDocument(); expect(screen.getByText("0.20")).toBeInTheDocument(); // The score did not decide this tier, so NO band explanation may render at all. // Asserting the absence of one specific band would pass vacuously: 0.20 sits in @@ -188,7 +190,9 @@ describe("RoutingDecisionCard", () => { // `signals` is gone under redaction; the cause alone must suppress the band. render(); expect(screen.queryByText(/SIMPLE|MEDIUM|COMPLEX|at or above/)).not.toBeInTheDocument(); - expect(screen.getByText("Heuristic, REASONING override (2 or more reasoning markers)")).toBeInTheDocument(); + expect( + screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score above the lowest tier)"), + ).toBeInTheDocument(); }); it("shows the operator's tier name on the badge instead of the canonical one", () => { @@ -212,7 +216,9 @@ describe("RoutingDecisionCard", () => { render( , ); - expect(screen.getByText("Heuristic, Deep override (2 or more reasoning markers)")).toBeInTheDocument(); + expect( + screen.getByText("Heuristic, Deep override (2 or more reasoning markers, score above the lowest tier)"), + ).toBeInTheDocument(); }); it("falls back to the raw cause for a value this build does not know", () => { 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 0876a539653..cb680668849 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -72,7 +72,7 @@ function describeCause(decision: RoutingDecision): string { case "heuristic_scorer": return "Heuristic scorer"; case "reasoning_override": - return `Heuristic, ${tierLabel ?? "REASONING"} override (2 or more reasoning markers)`; + return `Heuristic, ${tierLabel ?? "REASONING"} override (2 or more reasoning markers, score above the lowest tier)`; case "llm_classifier": return classifierModel ? `LLM classifier (${classifierModel})` : "LLM classifier"; case "literal_keyword_match": From b3c3e6ebb8a318809e13827389c5b9aada956e7c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:25:34 -0700 Subject: [PATCH 35/88] fix(search): default the AgentCore MCP protocol version to the gateway default --- litellm/llms/bedrock/search/transformation.py | 12 +++++++---- .../test_agentcore_search_transformation.py | 21 +++++++++++++++---- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 80e20d00ff9..94b69912faa 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -67,9 +67,12 @@ AGENTCORE_DEFAULT_TOOL_NAME: Final = "web-search-tool___WebSearch" AGENTCORE_TOOL_NAME_SUFFIX: Final = "___WebSearch" # MCP revision this provider speaks. Sent on every request because the gateway is -# called statelessly, without an initialize handshake to negotiate a version; -# servers that predate the header ignore it. -AGENTCORE_MCP_PROTOCOL_VERSION: Final = "2025-06-18" +# called statelessly, without an initialize handshake to negotiate a version. +# AgentCore gateways whose protocolConfiguration leaves supportedVersions unset +# accept only 2025-03-26 and reject anything newer with a -32600 error, so that +# is the default; a gateway pinned to another version needs +# AGENTCORE_MCP_PROTOCOL_VERSION set to match. +AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION: Final = "2025-03-26" # Matched against the URL host so a crafted path or query string can't pass for # a gateway hostname. @@ -168,7 +171,8 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): **headers, "Content-Type": "application/json", "Accept": "application/json, text/event-stream", - "MCP-Protocol-Version": AGENTCORE_MCP_PROTOCOL_VERSION, + "MCP-Protocol-Version": get_secret_str("AGENTCORE_MCP_PROTOCOL_VERSION") + or AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION, } def get_complete_url( diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 889f78c58b9..78808542c82 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -14,7 +14,7 @@ from unittest.mock import AsyncMock, patch, MagicMock import litellm from litellm.llms.bedrock.search.transformation import ( - AGENTCORE_MCP_PROTOCOL_VERSION, + AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION, AgentCoreSearchConfig, ) @@ -157,7 +157,20 @@ class TestAgentCoreSearch: headers = config.validate_environment(headers={}) assert headers["Accept"] == "application/json, text/event-stream" assert headers["Content-Type"] == "application/json" - assert headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION + assert headers["MCP-Protocol-Version"] == AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION + + def test_default_protocol_version_is_the_agentcore_gateway_default(self): + """A default AgentCore gateway supports only 2025-03-26 and answers + -32600 to anything newer, so that exact revision must be the default.""" + assert AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION == "2025-03-26" + + def test_protocol_version_env_override_wins(self): + """A gateway pinned to a newer supportedVersions list needs the header + to match, so AGENTCORE_MCP_PROTOCOL_VERSION must override the default.""" + config = AgentCoreSearchConfig() + with patch.dict(os.environ, {"AGENTCORE_MCP_PROTOCOL_VERSION": "2025-06-18"}): + headers = config.validate_environment(headers={}) + assert headers["MCP-Protocol-Version"] == "2025-06-18" def test_protocol_version_header_survives_signing(self): """Both auth paths must keep the MCP-Protocol-Version header on the wire.""" @@ -171,7 +184,7 @@ class TestAgentCoreSearch: api_base=GATEWAY_URL, api_key="test-jwt-token", ) - assert bearer_headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION + assert bearer_headers["MCP-Protocol-Version"] == AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION with patch.dict( os.environ, @@ -187,7 +200,7 @@ class TestAgentCoreSearch: api_base=GATEWAY_URL, ) assert signed_headers["Authorization"].startswith("AWS4-HMAC-SHA256") - assert signed_headers["MCP-Protocol-Version"] == AGENTCORE_MCP_PROTOCOL_VERSION + assert signed_headers["MCP-Protocol-Version"] == AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION def test_transform_search_response_parses_sse_frame(self): """Gateway may answer with an SSE-framed JSON-RPC message.""" From 20a3a16c2f6ba7a7d5755cf56740d5180973ede0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:25:43 -0700 Subject: [PATCH 36/88] fix(proxy): populate deployment fields on failed-request spend logs from the standard logging payload --- .../spend_tracking/spend_tracking_utils.py | 30 ++++++-- .../test_spend_tracking_utils.py | 77 +++++++++++++++++++ 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3146d8bccfb..0b56f0d8246 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -216,6 +216,15 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} +def _sl_attribution_fallback( + standard_logging_payload: StandardLoggingPayload | None, + field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"], +) -> str: + if standard_logging_payload is None: + return "" + return standard_logging_payload.get(field) or "" + + def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload: if kwargs is None: kwargs = {} @@ -288,8 +297,15 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs ): # use 'tags' from standard logging payload instead request_tags = safe_dumps(standard_logging_payload["request_tags"]) - _model_id: Final = metadata.get("model_info", {}).get("id", "") - _model_group: Final = metadata.get("model_group", "") + _model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback( + standard_logging_payload, "model_id" + ) + _model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback( + standard_logging_payload, "model_group" + ) + _api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback( + standard_logging_payload, "api_base" + ) # Extract overhead from hidden_params if available litellm_overhead_time_ms = None @@ -389,7 +405,11 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id") - custom_llm_provider: Final = kwargs.get("custom_llm_provider") + custom_llm_provider: Final = ( + kwargs.get("custom_llm_provider") + or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider") + or None + ) raw_model: Final = cast(str, kwargs.get("model") or "") model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) @@ -414,13 +434,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens), request_tags=request_tags, end_user=end_user_id or "", - api_base=litellm_params.get("api_base", ""), + api_base=_api_base, model_group=_model_group, model_id=_model_id, mcp_namespaced_tool_name=mcp_namespaced_tool_name, agent_id=agent_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), - custom_llm_provider=kwargs.get("custom_llm_provider", ""), + custom_llm_provider=custom_llm_provider or "", messages=_get_messages_for_spend_logs_payload( standard_logging_payload=standard_logging_payload, metadata=metadata ), diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e5add059260..9710dc44e99 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3164,3 +3164,80 @@ def test_batch_cost_row_id_is_stable_across_repeated_accounting(): ] assert ids[0] == ids[1] == "batch_same_batch_cost" + + +def _make_failed_request_standard_logging_payload() -> StandardLoggingPayload: + base: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + return cast( + StandardLoggingPayload, + { + **base, + "status": "failure", + "call_type": "aresponses", + "model_id": "mid-123", + "model_group": "group-x", + "api_base": "https://api.openai.com/v1/responses", + "custom_llm_provider": "openai", + }, + ) + + +def test_get_logging_payload_failed_request_falls_back_to_standard_logging_payload(): + """Failed-request kwargs from the proxy failure hook carry no deployment info + (LIT-5795), so the attribution columns must come from the failure-time + standard_logging_object.""" + payload = get_logging_payload( + kwargs={ + "model": "group-x", + "litellm_params": {"metadata": {"user_api_key": "test-key", "status": "failure"}}, + "standard_logging_object": _make_failed_request_standard_logging_payload(), + }, + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model_id"] == "mid-123" + assert payload["model_group"] == "group-x" + assert payload["api_base"] == "https://api.openai.com/v1/responses" + assert payload["custom_llm_provider"] == "openai" + + +def test_get_logging_payload_request_kwargs_win_over_standard_logging_payload(): + payload = get_logging_payload( + kwargs={ + "model": "group-y", + "custom_llm_provider": "anthropic", + "litellm_params": { + "api_base": "https://kwargs.example.com", + "metadata": { + "user_api_key": "test-key", + "model_group": "kwargs-group", + "model_info": {"id": "kwargs-mid"}, + }, + }, + "standard_logging_object": _make_failed_request_standard_logging_payload(), + }, + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model_id"] == "kwargs-mid" + assert payload["model_group"] == "kwargs-group" + assert payload["api_base"] == "https://kwargs.example.com" + assert payload["custom_llm_provider"] == "anthropic" + + +def test_get_logging_payload_failed_request_without_standard_logging_payload_leaves_fields_empty(): + payload = get_logging_payload( + kwargs={ + "model": "group-x", + "litellm_params": {"metadata": {"user_api_key": "test-key", "status": "failure"}}, + }, + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model_id"] == "" + assert payload["model_group"] == "" + assert payload["api_base"] == "" + assert payload["custom_llm_provider"] == "" From 0de60a2ff2d97ba31c35f5bde6aa08a31c11eae2 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 19 Aug 2026 14:26:02 -0700 Subject: [PATCH 37/88] fix(mcp): stop reporting failed OpenAPI tool calls as successes (#37496) An OpenAPI-backed MCP tool whose upstream answered 401 came back as a successful tool result carrying the upstream's rejection as its content, so a caller saw {"error":"invalid_token"} presented as data and the gateway recorded the request in its own spend log as call_mcp_tool | success. Three layers each erased the outcome. The request function returned response.text whatever the status, _handle_local_mcp_tool caught every exception and returned it as ordinary TextContent, and both dispatch sites then stamped isError=False unconditionally. Fixing only the first, which is the obvious fix, changes nothing, because the two above it still map failure onto the success-shaped value. The status is now classified where the response is held: a 401 becomes MCPUpstreamAuthError so the caller is told to re-authenticate, and every other non-2xx becomes MCPOpenApiUpstreamError, which carries the status and drops the upstream body rather than serving it as tool content. _handle_local_mcp_tool no longer swallows, and the call_tool arm keeps the auth error's type. Nothing new renders these: call_mcp_tool and call_tool_rest_api already turn them into an isError result naming the status and into a real 401 with WWW-Authenticate, and the OpenAPI path simply never reached them. The result is now byte-identical to the regular MCP path for the same failure. --- .../_experimental/mcp_server/exceptions.py | 18 +++ .../mcp_server/mcp_server_manager.py | 16 ++- .../mcp_server/openapi_to_mcp_generator.py | 64 +++++++-- .../proxy/_experimental/mcp_server/server.py | 32 +++-- .../mcp_server/test_mcp_server_manager.py | 69 +++++++++- .../test_openapi_to_mcp_generator.py | 125 +++++++++++++++++- .../mcp_server/test_openapi_tool_auth.py | 77 +++++++++++ 7 files changed, 370 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index 8c704c0fe93..a1b3b167a4a 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -75,6 +75,24 @@ class MCPUpstreamAuthError(Exception): ) +class MCPOpenApiUpstreamError(Exception): + """An OpenAPI-backed MCP tool's upstream answered with a non-2xx that is not a 401. + + Carries the status only. The upstream's response body is deliberately dropped rather than served + as tool content: it crosses a trust boundary and may hold prose, urls, or an error document that + reads as data, which is how these failures came to be reported as successful tool output. This + matches ``outcome_wire_value``'s contract for listing faults, category and status and nothing + else. A 401 is raised as ``MCPUpstreamAuthError`` instead, so the caller learns to + re-authenticate; every other status stays here, mirroring the regular MCP path where a 403 + deliberately does not produce a challenge. + """ + + def __init__(self, status_code: int, server_name: str) -> None: + self.status_code = status_code + self.server_name = server_name + super().__init__(f"upstream returned HTTP {status_code}") + + class MCPToolResultError(Exception): """An MCP tool call completed with ``isError=True`` in its result. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 65855df89f6..26a6f8d1251 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2330,7 +2330,15 @@ class MCPServerManager: input_schema = build_input_schema(resolved_operation) # Create tool function with headers using imported function - tool_func = create_tool_function(path, method, resolved_operation, base_url, headers=headers) + tool_func = create_tool_function( + path, + method, + resolved_operation, + base_url, + headers=headers, + server_label=server.name or server.server_name or server.alias or server.server_id, + relays_upstream_auth=server.is_client_forwarded_token, + ) tool_func.__name__ = prefixed_tool_name tool_func.__doc__ = description @@ -4979,6 +4987,12 @@ class MCPServerManager: return result + except MCPUpstreamAuthError: + # The caller must re-authenticate upstream, so this keeps its type all the way to the + # renderers: the streamable path turns it into an isError result naming the status, and + # the REST path relays a real 401 with the upstream's WWW-Authenticate. Flattening it + # into the generic message below would lose both. + raise except Exception as e: error_msg = f"Error calling OpenAPI tool {tool_name}: {e}" verbose_logger.error(error_msg) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index eb78aaeca0b..083a98cdd36 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -15,6 +15,12 @@ from urllib.parse import quote import httpx from typing_extensions import ReadOnly, Required +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPOpenApiUpstreamError, + MCPUpstreamAuthError, +) + # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to # ^[a-zA-Z0-9_-]+$ on tool names. Many specs (notably GitHub's REST API) use @@ -392,12 +398,40 @@ def _merge_openapi_tool_request_headers( return effective_headers +def _raise_for_upstream_failure( + response: httpx.Response, + upstream: str, + relays_upstream_auth: bool, +) -> None: + """Turn a non-2xx upstream response into the right typed failure, or return for a 2xx. + + Both call sites feed this: ``get`` hands back the response for a 4xx, while post/put/patch/delete + raise ``MaskedHTTPStatusError`` from inside the HTTP handler, so without one classifier the + non-GET tools would keep serving an error body as tool output. + + Only the client-forwarded modes carry the caller's own upstream token, so only they can act on a + 401 by re-authenticating; ``_call_regular_mcp_tool`` gates its re-auth signal the same way. Every + other status carries the code alone, never the upstream's body, which crosses a trust boundary. + """ + if response.status_code < 400: + return + if response.status_code == 401 and relays_upstream_auth: + raise MCPUpstreamAuthError( + status_code=response.status_code, + www_authenticate=response.headers.get("www-authenticate"), + server_name=upstream, + ) + raise MCPOpenApiUpstreamError(response.status_code, upstream) + + def create_tool_function( path: str, method: str, operation: _OpenAPIOperation, base_url: str, headers: dict[str, str] | None = None, + server_label: str | None = None, + relays_upstream_auth: bool = False, ): """Create a tool function for an OpenAPI operation. @@ -477,20 +511,26 @@ def create_tool_function( json_body = {"data": body_value} client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) + upstream: Final = server_label or f"{original_method.upper()} {path}" - if original_method == "get": - response = await client.get(url, params=params, headers=effective_headers) - elif original_method == "post": - response = await client.post(url, params=params, json=json_body, headers=effective_headers) - elif original_method == "put": - response = await client.put(url, params=params, json=json_body, headers=effective_headers) - elif original_method == "delete": - response = await client.delete(url, params=params, headers=effective_headers) - elif original_method == "patch": - response = await client.patch(url, params=params, json=json_body, headers=effective_headers) - else: - return f"Unsupported HTTP method: {original_method}" + try: + if original_method == "get": + response = await client.get(url, params=params, headers=effective_headers) + elif original_method == "post": + response = await client.post(url, params=params, json=json_body, headers=effective_headers) + elif original_method == "put": + response = await client.put(url, params=params, json=json_body, headers=effective_headers) + elif original_method == "delete": + response = await client.delete(url, params=params, headers=effective_headers) + elif original_method == "patch": + response = await client.patch(url, params=params, json=json_body, headers=effective_headers) + else: + return f"Unsupported HTTP method: {original_method}" + except MaskedHTTPStatusError as e: + _raise_for_upstream_failure(e.response, upstream, relays_upstream_auth) + raise + _raise_for_upstream_failure(response, upstream, relays_upstream_auth) return response.text return tool_function diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 8d69d84e492..0dc85c0318c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -407,8 +407,6 @@ if MCP_AVAILABLE: StreamableHTTPSessionManager = None from mcp.types import ( CallToolResult, - EmbeddedResource, - ImageContent, ListToolsResult, Prompt, TextContent, @@ -2861,12 +2859,11 @@ if MCP_AVAILABLE: _extra_token: Final = _request_extra_headers.set(forwarded_headers) _resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) try: - local_content = await _handle_local_mcp_tool(name, arguments) + response = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2940,8 +2937,7 @@ if MCP_AVAILABLE: if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args - local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=local_content, isError=False) + response = await _handle_local_mcp_tool(original_tool_name, arguments) return await _run_post_mcp_call_guardrails( result=response, @@ -3319,11 +3315,18 @@ if MCP_AVAILABLE: verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result - async def _handle_local_mcp_tool( - name: str, arguments: dict[str, object] - ) -> list[TextContent | ImageContent | EmbeddedResource]: - """ - Handle tool execution for local registry tools + async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult: + """Execute a local-registry tool and report whether it succeeded. + + Returns the result rather than bare content because the verdict is part of it: the content + alone cannot say whether the handler failed, so callers used to stamp isError=False on every + outcome and an upstream rejection was served as tool output. + + A failure is reported as ``isError=True`` here rather than raised, because the REST surface + turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash. + ``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to + re-authenticate, which both renderers already know how to say. + Note: Local tools don't use prefixes, so we use the original name """ import inspect @@ -3333,15 +3336,16 @@ if MCP_AVAILABLE: raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") try: - # Check if handler is async or sync if inspect.iscoroutinefunction(tool.handler): result = await tool.handler(**arguments) else: result = tool.handler(**arguments) - return [TextContent(text=str(result), type="text")] + except MCPUpstreamAuthError: + raise except Exception as e: verbose_logger.exception("Error executing local tool %s: %s", name, e) - return [TextContent(text=f"Error: {e}", type="text")] + return CallToolResult(content=[TextContent(text=f"Error: {e}", type="text")], isError=True) + return CallToolResult(content=[TextContent(text=str(result), type="text")], isError=False) def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9064312fd6d..5ee8143fb8e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -4395,8 +4395,12 @@ class TestMCPServerManager: captured: dict = {} - def fake_create_tool_function(path, method, operation, base_url, headers=None): + def fake_create_tool_function( + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + ): captured["headers"] = headers + captured["server_label"] = server_label + captured["relays_upstream_auth"] = relays_upstream_auth async def tool_func(**kwargs): return "ok" @@ -4425,6 +4429,11 @@ class TestMCPServerManager: assert captured["headers"] is not None assert captured["headers"]["Authorization"] == "STATIC token" + # The label names the server in an upstream-failure error, so registration must thread it; + # without this the fake would simply tolerate the argument and prove nothing about it. + assert captured["server_label"] == "openapi-server" + # auth_type is none here, so a 401 from this upstream must not be dressed up as a re-auth signal + assert captured["relays_upstream_auth"] is False @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): @@ -10765,3 +10774,61 @@ class TestResolveOpenapiToolAuth: ) assert "Authorization" not in (forwarded or {}) + + +class TestOpenApiHandlerRelaysUpstreamAuth: + """`_call_openapi_tool_handler` must not flatten a re-auth signal into a generic message. + + Its catch-all turned every exception into "Error calling OpenAPI tool ...", which is an isError + result but loses the status, so the REST surface could no longer relay a 401 with the upstream's + WWW-Authenticate and the streamable surface could not name the status the caller must act on. + """ + + @staticmethod + def _server() -> MCPServer: + return MCPServer( + server_id="srv-openapi", + name="report_api", + server_name="report_api", + url="https://api.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth_delegate, + spec_path="https://api.example.com/openapi.json", + ) + + @pytest.mark.asyncio + async def test_upstream_auth_error_keeps_its_type(self): + from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + manager = MCPServerManager() + server = self._server() + + async def raising_handler(**_kwargs): + raise MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer realm=x", server_name="report_api") + + tool = MagicMock() + tool.handler = raising_handler + with patch.object(global_mcp_tool_registry, "get_tool", return_value=tool): + with pytest.raises(MCPUpstreamAuthError) as exc: + await manager._call_openapi_tool_handler(server, "list_reports", {}) + + assert exc.value.status_code == 401 + assert exc.value.www_authenticate == "Bearer realm=x" + + @pytest.mark.asyncio + async def test_other_failures_still_become_an_error_result(self): + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + manager = MCPServerManager() + + async def raising_handler(**_kwargs): + raise RuntimeError("upstream returned HTTP 503") + + tool = MagicMock() + tool.handler = raising_handler + with patch.object(global_mcp_tool_registry, "get_tool", return_value=tool): + result = await manager._call_openapi_tool_handler(self._server(), "list_reports", {}) + + assert result.isError is True + assert "upstream returned HTTP 503" in result.content[0].text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 1f9316ee9c8..e59616e53c1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -27,12 +27,21 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( resolve_operation_params, ) +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPOpenApiUpstreamError, + MCPUpstreamAuthError, +) + GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" -def _create_mock_client(method: str, response_text: str) -> AsyncMock: - """Utility to create a mocked async httpx client for the given method.""" - response = SimpleNamespace(text=response_text) +def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock: + """Utility to create a mocked async httpx client for the given method. + + ``status_code`` and ``headers`` are part of the real response the tool function reads, so the + fake carries them too; a fake that omits them cannot observe whether the status is checked. + """ + response = SimpleNamespace(text=response_text, status_code=status_code, headers={}) client = AsyncMock() setattr(client, method, AsyncMock(return_value=response)) return client @@ -1259,3 +1268,113 @@ class TestRequestExtraHeaders: headers_sent = async_client.get.call_args[1]["headers"] assert "Authorization" not in headers_sent + + +class TestUpstreamStatusIsClassified: + """A non-2xx upstream must never be returned as tool output. + + The body used to be returned verbatim whatever the status, so an upstream rejection arrived as a + successful tool result and the request logged as a success. 401 is singled out because it is the + only status the caller can act on by re-authenticating, matching `_call_regular_mcp_tool` where a + 403 deliberately does not produce a challenge. + """ + + @staticmethod + def _tool(status_code: int, text: str = "body", headers: dict | None = None, relays_upstream_auth: bool = True): + response = SimpleNamespace(text=text, status_code=status_code, headers=headers or {}) + client = AsyncMock() + client.get = AsyncMock(return_value=response) + return create_tool_function( + "/reports", + "get", + {"operationId": "list_reports"}, + "https://api.example.com", + server_label="report_api", + relays_upstream_auth=relays_upstream_auth, + ), client + + @pytest.mark.asyncio + async def test_success_still_returns_the_body(self): + tool, client = self._tool(200, text='{"reports": []}') + with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): + assert await tool() == '{"reports": []}' + + @pytest.mark.asyncio + async def test_401_raises_the_reauth_signal_carrying_the_challenge(self): + tool, client = self._tool(401, text='{"error":"invalid_token"}', headers={"www-authenticate": 'Bearer realm="x"'}) + with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): + with pytest.raises(MCPUpstreamAuthError) as exc: + await tool() + + assert exc.value.status_code == 401 + assert exc.value.www_authenticate == 'Bearer realm="x"' + assert exc.value.server_name == "report_api" + + @pytest.mark.asyncio + async def test_401_on_a_non_forwarding_server_is_not_a_reauth_signal(self): + """Only the client-forwarded modes carry the caller's own upstream token, so only they can act + on a 401. `_call_regular_mcp_tool` gates its signal the same way, and without the gate an + api_key server rejecting a token would push clients into an OAuth flow that does not apply.""" + tool, client = self._tool(401, text='{"error":"bad key"}', relays_upstream_auth=False) + with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): + with pytest.raises(MCPOpenApiUpstreamError) as exc: + await tool() + + assert exc.value.status_code == 401 + + @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) + @pytest.mark.parametrize("status_code, expected", [(401, "auth"), (500, "other")]) + @pytest.mark.asyncio + async def test_non_get_methods_are_classified_too(self, method: str, status_code: int, expected: str): + """post/put/patch/delete call raise_for_status inside the HTTP handler. + + Only `get` hands a 4xx back to the caller; the others raise `MaskedHTTPStatusError` before any + status check the tool function could do, so classifying the returned response alone would + leave every non-GET tool still serving an upstream error body as successful tool output. A + fake client that simply returns a response cannot observe this, which is why this test builds + the error the real handler raises. + """ + import httpx + + from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError + + request = httpx.Request(method.upper(), "https://api.example.com/reports") + raw = httpx.Response( + status_code, + headers={"www-authenticate": 'Bearer realm="x"'}, + text="internal hostname db-prod-7.corp.example.com", + request=request, + ) + masked = MaskedHTTPStatusError(httpx.HTTPStatusError("boom", request=request, response=raw)) + + client = AsyncMock() + setattr(client, method, AsyncMock(side_effect=masked)) + tool = create_tool_function( + "/reports", + method, + {"operationId": "list_reports"}, + "https://api.example.com", + server_label="report_api", + relays_upstream_auth=True, + ) + + expected_type = MCPUpstreamAuthError if expected == "auth" else MCPOpenApiUpstreamError + with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): + with pytest.raises(expected_type) as exc: + await tool() + + assert exc.value.status_code == status_code + assert "db-prod-7" not in str(exc.value) + + @pytest.mark.parametrize("status_code", [403, 404, 429, 500, 503]) + @pytest.mark.asyncio + async def test_other_failures_raise_without_leaking_the_upstream_body(self, status_code: int): + secret_body = "internal hostname db-prod-7.corp.example.com and a stack trace" + tool, client = self._tool(status_code, text=secret_body) + with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): + with pytest.raises(MCPOpenApiUpstreamError) as exc: + await tool() + + assert exc.value.status_code == status_code + assert secret_body not in str(exc.value) + assert str(exc.value) == f"upstream returned HTTP {status_code}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 7bd846aeda4..bd953dc55f3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -652,3 +652,80 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc assert captured["resolver_credential"] == {"Authorization": OPENAPI_PER_SERVER_TOKEN} assert captured["injected"] == OPENAPI_PER_SERVER_TOKEN assert _request_auth_header.get() is None + + +@pytest.mark.parametrize("failure", ["auth", "other"]) +@pytest.mark.asyncio +async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: str): + """A failing local handler must never be reported as a successful tool result, and only an auth + failure may propagate. + + `_handle_local_mcp_tool` used to catch every exception and return it as TextContent, and both of + its callers then stamped `isError=False`, so an upstream rejection was served as tool output and + `extract_mcp_tool_result_error_message` logged the request as a success. + + The two kinds are split by consequence. `MCPUpstreamAuthError` propagates because both renderers + know it: the streamable path names the status and the REST path relays a real 401 with the + upstream's WWW-Authenticate. Anything else is reported as `isError=True` right here, because + `call_tool_rest_api` turns an unrecognized exception into HTTP 500 and an upstream 403 or 429 is + not a gateway crash. + """ + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPOpenApiUpstreamError, + MCPUpstreamAuthError, + ) + + error = ( + MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name="report_api") + if failure == "auth" + else MCPOpenApiUpstreamError(429, "report_api") + ) + + async def raising_handler(**_kwargs): + raise error + + fake_tool = MagicMock() + fake_tool.name = "list_reports" + fake_tool.handler = raising_handler + server = MCPServer( + server_id="srv-openapi", + name="report_api", + server_name="report_api", + url="https://api.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth_delegate, + spec_path="https://api.example.com/openapi.json", + ) + user = UserAPIKeyAuth(api_key="sk-user", user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) + + with ( + patch.object(mcp_module.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server), + patch.object(mcp_module.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={})), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool), + patch.object( + mcp_module.global_mcp_server_manager, + "resolve_openapi_upstream_auth", + new=AsyncMock(return_value=(None, None)), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + call = mcp_module.execute_mcp_tool( + name="list_reports", + arguments={}, + allowed_mcp_servers=[server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + if failure == "auth": + with pytest.raises(MCPUpstreamAuthError): + await call + return + result = await call + + # A non-auth upstream failure stays a 200 with isError, so REST does not report it as a gateway 500 + assert result.isError is True + assert "upstream returned HTTP 429" in result.content[0].text From 125587d286ed2a681bdb80e4e3bd661afa19fede Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:33:22 -0700 Subject: [PATCH 38/88] feat(e2e): canonical content-based match keys for record-and-replay Replay previously matched interactions by transport verb and path in recorded order, so a request whose body drifted from the recording silently replayed the stale response, and reordering two independent calls broke replay even though both were recorded. Match keys are now canonical: fixture_canonical.py strips volatile headers and credential fields, replaces unique markers, generated ids, uuids, and timestamps with fixed placeholders, sorts object keys, and hashes what remains, so a key is stable across runs and machines while any real content drift is a hard ReplayMiss naming the computed key, the closest recorded key with its file, and a content diff, with no fallthrough to a live call. Matching is order-independent across distinct keys and FIFO within one key. Recording now also redacts credential body and form fields (not just auth headers) so provider keys never land in bundles. Resolves LIT-5741 --- tests/e2e/CLAUDE.md | 6 +- tests/e2e/fixture_bundle.py | 19 +-- tests/e2e/fixture_canonical.py | 150 ++++++++++++++++++ tests/e2e/fixture_transport.py | 236 ++++++++++++++++++++++------ tests/e2e/test_fixture_canonical.py | 163 +++++++++++++++++++ tests/e2e/test_fixture_transport.py | 164 ++++++++++++++++++- 6 files changed, 672 insertions(+), 66 deletions(-) create mode 100644 tests/e2e/fixture_canonical.py create mode 100644 tests/e2e/test_fixture_canonical.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 9969ed10308..d7334552d0c 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -75,11 +75,11 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover `E2E_FIXTURE_MODE` selects the transport every client is built on: `live` (the default, and what an unset variable means: nothing changes), `record` (run against the live proxy and write every interaction to a fixture bundle), or `replay` (serve every interaction back from the bundle with no HTTP at all, so a replay run needs no proxy and cannot bill a provider). The seam is `select_transport` in `fixture_transport.py`, applied inside `build_proxy_client`; both transports fulfil the same `Transport` protocol, so no test or client changes shape in any mode -A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values and credential request fields (`api_key`, `*_secret_key`, `static_headers`, and the like; the list is `fixture_canonical.py`'s) are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format -Replay matches calls per test by transport verb and path in recorded order and raises `ReplayMiss` on any drift, naming the recorded and the actual call; a passed test must also consume its whole recording, or teardown fails it naming the first leftover interaction. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy +Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, path, and a content hash, so identity survives re-records and machine changes while any real content drift is a `ReplayMiss` that names the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a poll loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy -Deliberately not here yet: canonical content-based match keys (LIT-5741), streaming chunk fidelity (LIT-5742), and scoping record/replay to provider-bound traffic (LIT-5745) +Deliberately not here yet: streaming chunk fidelity (LIT-5742) and scoping record/replay to provider-bound traffic (LIT-5745) ## Typing diff --git a/tests/e2e/fixture_bundle.py b/tests/e2e/fixture_bundle.py index 5eff2cf2876..615ae8df1a4 100644 --- a/tests/e2e/fixture_bundle.py +++ b/tests/e2e/fixture_bundle.py @@ -8,10 +8,10 @@ green replay run can never certify against fixtures that have drifted more than a week from the live proxy. This module owns the format only. The transports that produce and consume it -live in fixture_transport.py; canonical request matching, streaming chunk -fidelity, and provider-scoping are follow-ups (LIT-5741/5742/5745) and are -deliberately absent here, which is why every interaction file stores the full -redacted request even though replay today matches by call order. +live in fixture_transport.py and the canonical match keys they compute live in +fixture_canonical.py (LIT-5741); streaming chunk fidelity and provider-scoping +are follow-ups (LIT-5742/5745). Every interaction file stores the full redacted +request because replay matches on its canonicalized content. """ from __future__ import annotations @@ -54,12 +54,13 @@ class Manifest(BaseModel): class RecordedRequest(BaseModel): - """The request as the transport saw it, auth header values redacted. + """The request as the transport saw it, auth header values and credential + body/form fields redacted. - Replay today only matches ``method`` (the transport verb, not the HTTP verb) - and ``path`` in call order; the rest is stored so LIT-5741 can move to - content-based match keys without re-recording. File uploads store a content - digest instead of the bytes.""" + Replay matches on the canonical content key fixture_canonical.py computes + over ``method`` (the transport verb, not the HTTP verb), ``path``, and the + canonicalized headers, params, body, form, and file identity. File uploads + store a content digest instead of the bytes.""" method: str path: str diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py new file mode 100644 index 00000000000..427f06bf8fb --- /dev/null +++ b/tests/e2e/fixture_canonical.py @@ -0,0 +1,150 @@ +"""Canonical request identity for replay matching (LIT-5741). + +Matching a replayed call against the raw recorded request never hits: unique +markers salt prompts, model names, and tags; every run mints fresh virtual +keys; request ids and timestamps differ on every call. Matching on transport +verb + path alone collides: two different requests to the same route silently +swap responses, which passes when it should miss. The canonicalizer strips +exactly the volatile material (volatile headers, credential fields, markers, +generated ids, timestamps) and hashes what remains with sorted object keys, so +identity is content-based and stable across runs and machines. + +Every rewrite rule lives in this module, next to the transports that apply it: +a new volatile header, credential field name, or generated-id shape is one +edit here, never a per-suite change. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass +from functools import reduce +from typing import Final + +from pydantic import JsonValue + +from fixture_bundle import RecordedRequest + +VOLATILE_HEADER_NAMES: Final[frozenset[str]] = frozenset( + { + "authorization", + "x-litellm-api-key", + "x-api-key", + "x-goog-api-key", + "x-request-id", + "traceparent", + "tracestate", + } +) + +SECRET_FIELD_NAMES: Final[frozenset[str]] = frozenset( + {"api_key", "aws_access_key_id", "static_headers", "vertex_credentials"} +) +SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = ( + "_api_key", + "_secret_key", + "_secret_access_key", + "_session_token", + "_credentials", + "_password", +) +SECRET_PLACEHOLDER: Final = "" + +PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( + (re.compile(r"(?"), + ( + re.compile(r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"), + "", + ), + (re.compile(r"sk-[A-Za-z0-9_-]{16,}"), ""), + ( + re.compile(r"\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?"), + "", + ), + (re.compile(r"(?"), + ( + re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"), + "", + ), + (re.compile(r"(?"), +) + + +def is_secret_field(name: str) -> bool: + lowered: Final = name.lower() + return lowered in SECRET_FIELD_NAMES or lowered.endswith(SECRET_FIELD_SUFFIXES) + + +def canonical_string(value: str) -> str: + return reduce(lambda acc, rule: rule[0].sub(rule[1], acc), PLACEHOLDER_RULES, value) + + +def _canonical_flat(fields: dict[str, str]) -> dict[str, JsonValue]: + return { + key: SECRET_PLACEHOLDER if is_secret_field(key) else canonical_string(value) + for key, value in fields.items() + } + + +def _canonical_value(value: JsonValue) -> JsonValue: + match value: + case str(): + return canonical_string(value) + case dict(): + return { + key: SECRET_PLACEHOLDER + if is_secret_field(key) and item is not None + else _canonical_value(item) + for key, item in value.items() + } + case list(): + return [_canonical_value(item) for item in value] + case _: + return value + + +@dataclass(frozen=True, slots=True) +class CanonicalRequest: + method: str + path: str + content: str + + @property + def key(self) -> str: + digest: Final = hashlib.sha256( + f"{self.method} {self.path}\n{self.content}".encode() + ).hexdigest()[:16] + return f"{self.method} {self.path} #{digest}" + + def pretty_content(self) -> str: + return json.dumps(json.loads(self.content), indent=2, sort_keys=True) + + +def canonicalize(request: RecordedRequest) -> CanonicalRequest: + file_identity: Final[JsonValue | None] = ( + None + if request.file_name is None and request.file_sha256 is None + else { + "name": None if request.file_name is None else canonical_string(request.file_name), + "sha256": request.file_sha256, + "bytes": request.file_bytes, + } + ) + content: Final[dict[str, JsonValue]] = { + "headers": { + name.lower(): canonical_string(value) + for name, value in request.headers.items() + if name.lower() not in VOLATILE_HEADER_NAMES + }, + "params": _canonical_flat(request.params), + "body": _canonical_value(request.body), + "form": None if request.form is None else _canonical_flat(request.form), + "file": file_identity, + } + return CanonicalRequest( + method=request.method, + path=canonical_string(request.path), + content=json.dumps(content, sort_keys=True, separators=(",", ":")), + ) diff --git a/tests/e2e/fixture_transport.py b/tests/e2e/fixture_transport.py index b99b0d2d80d..d49ac63787a 100644 --- a/tests/e2e/fixture_transport.py +++ b/tests/e2e/fixture_transport.py @@ -7,23 +7,30 @@ HTTP, no proxy, no provider spend. Because both fulfil ``Transport``, no test or client changes shape; ``build_proxy_client`` picks the transport from ``E2E_FIXTURE_MODE`` (live | record | replay, default live). -Replay matches each call by test node id and call order, verifying transport -verb + path and failing hard on any drift (``ReplayMiss``). Canonical -content-based match keys are LIT-5741; streaming chunk fidelity is LIT-5742; -scoping record/replay to provider-bound traffic is LIT-5745. +Replay matches each call by test node id and canonical content key +(fixture_canonical.py, LIT-5741): volatile headers, credential fields, unique +markers, generated ids, and timestamps are canonicalized out before hashing, so +matching is order-independent across distinct keys, FIFO within a key, and a +miss fails hard (``ReplayMiss``) printing the computed key and the closest +recorded key without ever falling through to a live call. Streaming chunk +fidelity is LIT-5742; scoping record/replay to provider-bound traffic is +LIT-5745. """ from __future__ import annotations +import difflib import functools import hashlib import os +from collections import deque from dataclasses import dataclass, field from datetime import datetime +from itertools import islice from pathlib import Path from typing import Final, Literal, assert_never -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue from e2e_http import AuthHeaders, BinaryStream, ProbeResult, Result, StreamingResponse from fixture_bundle import ( @@ -43,12 +50,14 @@ from fixture_bundle import ( check_freshness, format_age, from_result, + interaction_filename, load_bundle, prepare_bundle, slug_for_test, to_json_value, to_result, ) +from fixture_canonical import CanonicalRequest, canonicalize, is_secret_field from transport import Transport type FixtureMode = Literal["live", "record", "replay"] @@ -118,6 +127,27 @@ def _redact(headers: dict[str, str]) -> dict[str, str]: } +def _redact_secret_fields(value: JsonValue) -> JsonValue: + match value: + case dict(): + return { + key: REDACTED_VALUE + if is_secret_field(key) and item is not None + else _redact_secret_fields(item) + for key, item in value.items() + } + case list(): + return [_redact_secret_fields(item) for item in value] + case _: + return value + + +def _redact_flat(fields: dict[str, str]) -> dict[str, str]: + return { + key: REDACTED_VALUE if is_secret_field(key) else value for key, value in fields.items() + } + + def recorded_request( method: str, path: str, @@ -133,9 +163,9 @@ def recorded_request( method=method, path=path, headers=_redact(_dump_flat(headers)), - params=_dump_flat(params), - body=None if body is None else to_json_value(body), - form=None if form is None else _dump_flat(form), + params=_redact_flat(_dump_flat(params)), + body=None if body is None else _redact_secret_fields(to_json_value(body)), + form=None if form is None else _redact_flat(_dump_flat(form)), file_name=file_name, file_sha256=None if file_content is None else hashlib.sha256(file_content).hexdigest(), file_bytes=None if file_content is None else len(file_content), @@ -304,46 +334,106 @@ class RecordingTransport: return response +def _build_pool(recorded: tuple[Interaction, ...]) -> dict[str, deque[Interaction]]: + keys: Final = tuple(canonicalize(interaction.request).key for interaction in recorded) + return { + key: deque( + interaction + for candidate_key, interaction in zip(keys, recorded, strict=True) + if candidate_key == key + ) + for key in dict.fromkeys(keys) + } + + +def _closest_recorded( + canonical: CanonicalRequest, recorded: tuple[Interaction, ...] +) -> tuple[CanonicalRequest, str]: + candidates: Final = tuple(canonicalize(interaction.request) for interaction in recorded) + ratios: Final = tuple( + difflib.SequenceMatcher( + None, f"{canonical.method} {canonical.path}\n{canonical.content}", + f"{candidate.method} {candidate.path}\n{candidate.content}", + ).ratio() + for candidate in candidates + ) + best: Final = max(range(len(candidates)), key=lambda index: ratios[index]) + return candidates[best], interaction_filename(best, recorded[best].request) + + +def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle: LoadedBundle) -> str: + recorded: Final = bundle.interactions.get(slug, ()) + if not recorded: + return ( + f"replay miss for {test_key}: computed key {canonical.key} but nothing is recorded " + f"under {slug}; re-record with E2E_FIXTURE_MODE=record" + ) + closest, closest_file = _closest_recorded(canonical, recorded) + diff: Final = "\n".join( + islice( + difflib.unified_diff( + closest.pretty_content().splitlines(), + canonical.pretty_content().splitlines(), + fromfile=f"closest recorded ({closest_file})", + tofile="test made", + lineterm="", + ), + 60, + ) + ) + return ( + f"replay miss for {test_key}: no recorded interaction matches key {canonical.key}; " + f"closest recorded key is {closest.key} ({closest_file})\n{diff}\n" + "re-record with E2E_FIXTURE_MODE=record" + ) + + @dataclass(slots=True) class ReplaySource: - """One shared cursor set over a loaded bundle, so every client built in the - session consumes the same recorded sequence per test.""" + """One shared pool per test over a loaded bundle, so every client built in + the session consumes the same recorded interactions. Calls match by + canonical content key: order-independent across distinct keys (concurrent + tests interleave calls nondeterministically), FIFO within one key (a poll + loop replays its recorded responses in recorded order).""" bundle: LoadedBundle - _cursors: dict[str, int] = field(default_factory=dict) + _pools: dict[str, dict[str, deque[Interaction]]] = field(default_factory=dict) - def next_interaction(self, method: str, path: str) -> Interaction: - test_key = current_test_key() - slug = slug_for_test(test_key) - recorded = self.bundle.interactions.get(slug, ()) - index = self._cursors.get(slug, 0) - if index >= len(recorded): + def _pool(self, slug: str) -> dict[str, deque[Interaction]]: + if slug not in self._pools: + self._pools[slug] = _build_pool(self.bundle.interactions.get(slug, ())) + return self._pools[slug] + + def next_interaction(self, request: RecordedRequest) -> Interaction: + test_key: Final = current_test_key() + slug: Final = slug_for_test(test_key) + pool: Final = self._pool(slug) + canonical: Final = canonicalize(request) + queue: Final = pool.get(canonical.key) + if queue is None: + raise ReplayMiss(_miss_message(test_key, slug, canonical, self.bundle)) + if not queue: raise ReplayMiss( - f"replay exhausted for {test_key}: call #{index + 1} ({method} {path}) has no recorded " - f"interaction ({len(recorded)} recorded under {slug}); re-record with E2E_FIXTURE_MODE=record" + f"replay exhausted for {test_key}: every recorded interaction for key " + f"{canonical.key} is already consumed; re-record with E2E_FIXTURE_MODE=record" ) - interaction = recorded[index] - if interaction.request.method != method or interaction.request.path != path: - raise ReplayMiss( - f"replay mismatch for {test_key} at call #{index + 1}: recorded " - f"{interaction.request.method} {interaction.request.path}, test made {method} {path}; " - "re-record with E2E_FIXTURE_MODE=record" - ) - self._cursors[slug] = index + 1 - return interaction + return queue.popleft() def leftover_error(self, test_key: str) -> str | None: """Non-None when the test consumed fewer interactions than were recorded, meaning a passing replay proved less than the bundle claims.""" - slug = slug_for_test(test_key) - recorded = self.bundle.interactions.get(slug, ()) - consumed = self._cursors.get(slug, 0) - if consumed >= len(recorded): + slug: Final = slug_for_test(test_key) + recorded: Final = self.bundle.interactions.get(slug, ()) + if not recorded: + return None + leftover: Final = tuple( + interaction for queue in self._pool(slug).values() for interaction in queue + ) + if not leftover: return None - pending = recorded[consumed] return ( - f"replay incomplete for {test_key}: {len(recorded) - consumed} of {len(recorded)} recorded " - f"interactions never consumed, next is {pending.request.method} {pending.request.path}; " + f"replay incomplete for {test_key}: {len(leftover)} of {len(recorded)} recorded " + f"interactions never consumed, e.g. {canonicalize(leftover[0].request).key}; " "re-record with E2E_FIXTURE_MODE=record" ) @@ -386,7 +476,12 @@ class ReplayTransport: def post[R: BaseModel]( self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("post", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction(recorded_request("post", path, headers=headers, body=json)) + ), + response_type, + ) def get[R: BaseModel]( self, @@ -397,7 +492,12 @@ class ReplayTransport: response_type: type[R], timeout: float | None = None, ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("get", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction(recorded_request("get", path, headers=headers, params=params)) + ), + response_type, + ) def delete[R: BaseModel]( self, @@ -408,25 +508,46 @@ class ReplayTransport: response_type: type[R], params: BaseModel | None = None, ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("delete", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction( + recorded_request("delete", path, headers=headers, body=json, params=params) + ) + ), + response_type, + ) def patch[R: BaseModel]( self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("patch", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction(recorded_request("patch", path, headers=headers, body=json)) + ), + response_type, + ) def put[R: BaseModel]( self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("put", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction(recorded_request("put", path, headers=headers, body=json)) + ), + response_type, + ) def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: - return _expect_streaming(self.source.next_interaction("stream", path)) + return _expect_streaming( + self.source.next_interaction(recorded_request("stream", path, headers=headers, body=json)) + ) def stream_binary( self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 ) -> BinaryStream: - interaction = self.source.next_interaction("stream_binary", path) + interaction = self.source.next_interaction( + recorded_request("stream_binary", path, headers=headers, body=json) + ) match interaction.response: case RecordedBinary(payload=payload): return payload @@ -444,10 +565,16 @@ class ReplayTransport: params: BaseModel | None = None, stream: bool = False, ) -> StreamingResponse: - return _expect_streaming(self.source.next_interaction("send", path)) + return _expect_streaming( + self.source.next_interaction( + recorded_request("send", path, headers=headers, body=json, params=params) + ) + ) def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - interaction = self.source.next_interaction("probe", path) + interaction = self.source.next_interaction( + recorded_request("probe", path, headers=self.master, params=params) + ) match interaction.response: case RecordedProbe(payload=payload): return payload @@ -467,10 +594,27 @@ class ReplayTransport: params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: - return to_result(_expect_result(self.source.next_interaction("upload", path)), response_type) + return to_result( + _expect_result( + self.source.next_interaction( + recorded_request( + "upload", + path, + headers=headers, + params=params, + form=form, + file_name=filename, + file_content=content, + ) + ) + ), + response_type, + ) def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: - return _expect_streaming(self.source.next_interaction("download", path)) + return _expect_streaming( + self.source.next_interaction(recorded_request("download", path, headers=headers)) + ) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/test_fixture_canonical.py b/tests/e2e/test_fixture_canonical.py new file mode 100644 index 00000000000..30c57dc3ac6 --- /dev/null +++ b/tests/e2e/test_fixture_canonical.py @@ -0,0 +1,163 @@ +"""Harness coverage for canonical request identity (LIT-5741). + +No proxy and no ``e2e`` marker: pure functions over ``RecordedRequest``. Pins +the two failure modes match keys must avoid: keying on volatile material so +nothing ever matches (markers, virtual keys, ids, timestamps, volatile +headers), and keying on too little so different requests collide and a test +silently asserts against another request's response. +""" + +from __future__ import annotations + +import pytest +from pydantic import JsonValue + +from fixture_bundle import RecordedRequest +from fixture_canonical import CanonicalRequest, canonical_string, canonicalize, is_secret_field + + +def request( + method: str = "post", + path: str = "/chat/completions", + *, + headers: dict[str, str] | None = None, + params: dict[str, str] | None = None, + body: JsonValue | None = None, + form: dict[str, str] | None = None, + file_name: str | None = None, + file_sha256: str | None = None, + file_bytes: int | None = None, +) -> RecordedRequest: + return RecordedRequest( + method=method, + path=path, + headers=headers or {}, + params=params or {}, + body=body, + form=form, + file_name=file_name, + file_sha256=file_sha256, + file_bytes=file_bytes, + ) + + +class TestPlaceholders: + @pytest.mark.parametrize( + ("raw", "expected"), + [ + ("Reply ok. 4d5152a995b7", "Reply ok. "), + ("e2e-chat-stream-4d5152a995b7", "e2e-chat-stream-"), + ("sk-3mCXCTGmYuEEIU2i2qmVE3Xq6tSK1O0X6ZIRP1Lpw8ZlbNjt", ""), + ("9f1c8a2e-4b3d-4f6a-8f2f-0a1b2c3d4e5f", ""), + ("z" * 64, "z" * 64), + ("0123456789abcdef" * 4, ""), + ("2026-08-19T20:57:13.363499+00:00", ""), + ("2026-08-19", ""), + ("chatcmpl-C0LO6rRkfJlpJ2mqW9BHYo4Sm8FWl", ""), + ("batch_688a8b7f9a08819096e0f7c88fcd07c5", ""), + ("file-XyZ12345abc", ""), + ("gpt-4o-mini", "gpt-4o-mini"), + ("max_tokens", "max_tokens"), + ("sk-1234", "sk-1234"), + ], + ) + def test_rewrites_exactly_the_volatile_shapes(self, raw: str, expected: str) -> None: + assert canonical_string(raw) == expected + + +class TestSecretFields: + @pytest.mark.parametrize( + ("name", "secret"), + [ + ("api_key", True), + ("openai_api_key", True), + ("aws_secret_access_key", True), + ("aws_session_token", True), + ("vertex_credentials", True), + ("static_headers", True), + ("langfuse_secret_key", True), + ("model", False), + ("max_completion_tokens", False), + ("api_base", False), + ], + ) + def test_names_that_carry_credentials(self, name: str, secret: bool) -> None: + assert is_secret_field(name) is secret + + +class TestKeyStability: + def test_volatile_material_does_not_change_the_key(self) -> None: + """Acceptance: a suite recorded on one machine (fresh keys, that day's + dates, that run's markers) replays on another with no misses.""" + first = request( + headers={"authorization": "Bearer sk-run-one-aaaaaaaaaaaaaaaa", "x-request-id": "req-1"}, + params={"start_date": "2026-08-18"}, + body={ + "model": "e2e-chat-4d5152a995b7", + "messages": [{"role": "user", "content": "Reply ok. 4d5152a995b7"}], + "api_key": "sk-live-one-aaaaaaaaaaaaaaaa", + }, + ) + second = request( + headers={"authorization": "Bearer sk-run-two-bbbbbbbbbbbbbbbb", "x-request-id": "req-2"}, + params={"start_date": "2026-08-19"}, + body={ + "model": "e2e-chat-1a2b3c4d5e6f", + "messages": [{"role": "user", "content": "Reply ok. 1a2b3c4d5e6f"}], + "api_key": "os.environ/OPENAI_API_KEY", + }, + ) + assert canonicalize(first).key == canonicalize(second).key + + def test_serialization_order_is_not_identity(self) -> None: + ordered = request(body={"model": "m", "stream": True}) + reversed_order = request(body={"stream": True, "model": "m"}) + assert canonicalize(ordered).key == canonicalize(reversed_order).key + + def test_generated_ids_in_the_path_do_not_change_the_key(self) -> None: + first = request("get", "/v1/batches/batch_688a8b7f9a08819096e0f7c88fcd07c5") + second = request("get", "/v1/batches/batch_770b9c8f0b19920107f1f8d99fde18d6") + assert canonicalize(first).key == canonicalize(second).key + + +class TestKeyDistinctness: + def test_requests_differing_only_inside_canonicalized_fields_stay_distinct(self) -> None: + """Acceptance: a naive verb+path hash collides these; the content key + must not, or one test silently asserts against the other's response.""" + first = request(body={"messages": [{"content": "Reply ok. 4d5152a995b7"}]}) + second = request(body={"messages": [{"content": "Count to three. 4d5152a995b7"}]}) + naive = (first.method, first.path) + assert naive == (second.method, second.path) + assert canonicalize(first).key != canonicalize(second).key + + def test_a_kept_header_is_identity(self) -> None: + first = request(headers={"x-litellm-tags": "prod"}) + second = request(headers={"x-litellm-tags": "shadow"}) + assert canonicalize(first).key != canonicalize(second).key + + def test_a_volatile_header_is_not_identity(self) -> None: + first = request(headers={"traceparent": "00-aa-bb-01", "x-api-key": "one"}) + second = request(headers={"traceparent": "00-cc-dd-01", "x-api-key": "two"}) + assert canonicalize(first).key == canonicalize(second).key + + def test_secret_set_versus_unset_stays_distinct(self) -> None: + with_key = request(body={"api_key": "sk-live-aaaaaaaaaaaaaaaa"}) + without_key = request(body={"api_key": None}) + assert canonicalize(with_key).key != canonicalize(without_key).key + + def test_file_content_is_identity(self) -> None: + first = request( + "upload", "/v1/files", file_name="batch.jsonl", file_sha256="a" * 64, file_bytes=10 + ) + second = request( + "upload", "/v1/files", file_name="batch.jsonl", file_sha256="b" * 64, file_bytes=10 + ) + assert canonicalize(first).key != canonicalize(second).key + + +class TestKeyShape: + def test_key_names_method_path_and_digest(self) -> None: + canonical = canonicalize(request("post", "/model/new", body={"model_name": "m"})) + assert isinstance(canonical, CanonicalRequest) + assert canonical.key.startswith("post /model/new #") + assert len(canonical.key.rsplit("#", 1)[1]) == 16 diff --git a/tests/e2e/test_fixture_transport.py b/tests/e2e/test_fixture_transport.py index 5c7201cca37..814155f64d1 100644 --- a/tests/e2e/test_fixture_transport.py +++ b/tests/e2e/test_fixture_transport.py @@ -5,9 +5,10 @@ the live one (dependency injection, no monkeypatching): recording must pass every value through unchanged while writing one redacted interaction file per call, and replay must serve identical values from the bundle alone - the fake's call log proves nothing reaches the inner transport - failing hard -(``ReplayMiss``) on any drift in order, verb, or path. The collection-time -gate and report header are pinned here too, including the stale message that -names the bundle's age. +(``ReplayMiss``) on any content drift, printing the computed canonical key and +the closest recorded key (LIT-5741; the pure canonicalizer is pinned in +test_fixture_canonical.py). The collection-time gate and report header are +pinned here too, including the stale message that names the bundle's age. """ from __future__ import annotations @@ -16,6 +17,7 @@ import hashlib from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path +from uuid import uuid4 import pytest from pydantic import BaseModel @@ -35,10 +37,12 @@ from fixture_bundle import ( Interaction, LoadedBundle, Manifest, + RecordedResult, load_bundle, prepare_bundle, slug_for_test, ) +from fixture_canonical import canonicalize from fixture_transport import ( InvalidFixtureMode, RecordingTransport, @@ -50,6 +54,7 @@ from fixture_transport import ( fixture_mode_collection_error, fixture_report_lines, parse_fixture_mode, + recorded_request, replay_leftover_error, select_transport, ) @@ -70,6 +75,17 @@ class Query(BaseModel): q: str +class DeployParams(BaseModel): + model: str + api_key: str | None = None + aws_secret_access_key: str | None = None + + +class DeployBody(BaseModel): + model_name: str + litellm_params: DeployParams + + STREAMING = StreamingResponse( status_code=200, body="", @@ -272,6 +288,28 @@ class TestRecordingTransport: } assert "sk-secret" not in this_tests_files(root)[0].read_text(encoding="utf-8") + def test_redacts_credential_body_fields_in_the_recorded_request(self, tmp_path: Path) -> None: + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post( + "/model/new", + headers=fake.master, + json=DeployBody( + model_name="m", + litellm_params=DeployParams(model="openai/gpt", api_key="sk-live-provider-secret-123456"), + ), + response_type=Payload, + ) + raw = this_tests_files(root)[0].read_text(encoding="utf-8") + interaction = Interaction.model_validate_json(raw) + assert "sk-live-provider-secret-123456" not in raw + assert isinstance(interaction.request.body, dict) + params = interaction.request.body["litellm_params"] + assert isinstance(params, dict) + assert params["api_key"] == "" + assert params["aws_secret_access_key"] is None + def test_upload_records_a_content_digest_not_the_bytes(self, tmp_path: Path) -> None: fake = FakeTransport() root = tmp_path / "bundle" @@ -335,25 +373,135 @@ class TestReplayTransport: ) assert fake.calls == calls_after_record - def test_mismatched_call_names_recorded_and_actual(self, tmp_path: Path) -> None: + def test_miss_names_the_computed_key_and_the_closest_recorded_key(self, tmp_path: Path) -> None: fake = FakeTransport() root = tmp_path / "bundle" recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - with pytest.raises(ReplayMiss, match=r"recorded post /model/new, test made get /v1/models"): + with pytest.raises(ReplayMiss) as excinfo: replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) + message = str(excinfo.value) + assert "no recorded interaction matches key get /v1/models #" in message + assert "closest recorded key is post /model/new #" in message + assert "0000-post-model-new.json" in message + assert "re-record with E2E_FIXTURE_MODE=record" in message - def test_exhausted_recording_names_the_call_count(self, tmp_path: Path) -> None: + def test_content_drift_on_the_same_route_misses_with_no_live_call(self, tmp_path: Path) -> None: + """The naive verb+path match replayed a stale response for a request + whose content had changed, silently passing; a content key must miss, + print both canonical forms' diff, and never reach the inner transport.""" + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + calls_after_record = list(fake.calls) + replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") + with pytest.raises(ReplayMiss) as excinfo: + replay.post("/model/new", headers=replay.master, json=Body(prompt="y"), response_type=Payload) + message = str(excinfo.value) + assert "no recorded interaction matches key post /model/new #" in message + assert "closest recorded key is post /model/new #" in message + assert '- "prompt": "x"' in message + assert '+ "prompt": "y"' in message + assert fake.calls == calls_after_record + + def test_exhausted_key_names_the_key(self, tmp_path: Path) -> None: fake = FakeTransport() root = tmp_path / "bundle" recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - with pytest.raises(ReplayMiss, match=r"call #2 \(post /model/new\) has no recorded interaction \(1 recorded"): + with pytest.raises( + ReplayMiss, match=r"every recorded interaction for key post /model/new #\w{16} is already consumed" + ): replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + def test_replays_out_of_recorded_order_across_distinct_keys(self, tmp_path: Path) -> None: + """Concurrent tests interleave independent calls nondeterministically + (e.g. a burst of parallel chat calls), so replay matches by content, + never by recorded position.""" + fake = FakeTransport() + root = tmp_path / "bundle" + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) + recording.post("/key/generate", headers=fake.master, json=Body(prompt="k"), response_type=Payload) + source = replay_source(root) + replay: Transport = ReplayTransport(source=source, master_key="sk-1234") + replay.post("/key/generate", headers=replay.master, json=Body(prompt="k"), response_type=Payload) + replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) + assert source.leftover_error(current_test_key()) is None + + def test_identical_requests_replay_their_responses_in_recorded_order(self, tmp_path: Path) -> None: + """A poll loop makes the same request repeatedly and asserts on the + progression, so duplicates under one key stay FIFO.""" + root = tmp_path / "bundle" + recorder = make_recorder(root) + recorder.record( + test_key=current_test_key(), + request=recorded_request( + "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") + ), + response=RecordedResult(kind="success", status_code=200, data={"value": "first"}), + ) + recorder.record( + test_key=current_test_key(), + request=recorded_request( + "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") + ), + response=RecordedResult(kind="success", status_code=200, data={"value": "second"}), + ) + replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") + first = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) + second = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) + assert first == Success(status_code=200, data=Payload(value="first")) + assert second == Success(status_code=200, data=Payload(value="second")) + + +class TestRecordedKeySets: + def test_two_separate_recordings_of_one_flow_produce_identical_key_sets( + self, tmp_path: Path + ) -> None: + """Everything a run randomizes (markers, virtual keys, dates) must + canonicalize out, so separately recorded runs of the same suite agree + on every match key and a bundle recorded elsewhere replays here.""" + + def record_flow(root: Path, run_date: str) -> list[str]: + fake = FakeTransport() + recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) + marker = deterministic_marker() + recording.post( + "/model/new", + headers=fake.master, + json=DeployBody( + model_name=f"e2e-chat-{marker}", + litellm_params=DeployParams(model="openai/gpt", api_key=f"sk-live-{uuid4().hex}"), + ), + response_type=Payload, + ) + recording.post( + "/chat/completions", + headers=recording.bearer(f"sk-{uuid4().hex}"), + json=Body(prompt=f"Reply with the single word ok. {marker}"), + response_type=Payload, + ) + recording.get( + "/spend/logs", headers=fake.master, params=Query(q=run_date), response_type=Payload + ) + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + return sorted( + canonicalize(interaction.request).key + for interactions in loaded.interactions.values() + for interaction in interactions + ) + + first_keys = record_flow(tmp_path / "one", "2026-08-18") + second_keys = record_flow(tmp_path / "two", "2026-08-19") + assert first_keys == second_keys + assert len(first_keys) == 3 + class TestReplayLeftover: def test_fully_consumed_recording_leaves_nothing(self, tmp_path: Path) -> None: @@ -378,7 +526,7 @@ class TestReplayLeftover: error = source.leftover_error(current_test_key()) assert error is not None assert "1 of 2 recorded interactions never consumed" in error - assert "next is probe /health/liveliness" in error + assert "e.g. probe /health/liveliness #" in error assert "re-record with E2E_FIXTURE_MODE=record" in error def test_test_without_recordings_has_no_leftover(self, tmp_path: Path) -> None: From 5eeccf69b60b15e1f0939371d3ee18c0eb2099d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:35:01 -0700 Subject: [PATCH 39/88] fix(batches): skip undecodable batch output lines when costing --- litellm/batches/batch_utils.py | 2 +- tests/test_litellm/batches/test_batch_utils.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0f8acce3379..0cf22d82ca6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -426,7 +426,7 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]: def _parse_batch_output_line(line: bytes) -> dict | None: try: parsed: Final = json.loads(line) - except json.JSONDecodeError as e: + except ValueError as e: verbose_logger.warning("skipping malformed batch output line: %s", str(e)) return None if isinstance(parsed, dict): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index f3a7413bf03..ebe093c591c 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -186,6 +186,11 @@ def test_iter_output_entries_skips_malformed_and_non_object_lines(): assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}] +def test_iter_output_entries_skips_undecodable_line(): + content = b'{"ok": 1}\n{"note": "\xff-bad"}\n{"ok": 2}\n' + assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}] + + # =========================================================================== # # _estimate_batch_entry_tokens (regression: an uncountable/malformed row must # never contribute zero tokens, or a crafted batch could evade the TPM limit) From ed84e82428aa39903a5530998717a5afd252be4b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:35:25 -0700 Subject: [PATCH 40/88] docs(search): use the latest Bedrock sonnet in the agentcore example config --- .../proxy/example_config_yaml/agentcore_websearch_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml index f2c5a460bf0..bc29794e2a3 100644 --- a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml +++ b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml @@ -5,7 +5,7 @@ model_list: - model_name: claude-sonnet litellm_params: - model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 + model: bedrock/us.anthropic.claude-sonnet-5 aws_region_name: us-east-1 search_tools: From 8ef522a2a034239b8544ba3febec03883ad4346e Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 19 Aug 2026 21:41:57 +0000 Subject: [PATCH 41/88] fix(search): read AgentCore structuredContent results Web-search connector 1.1.0 and later return the machine-readable results in result.structuredContent and may leave the text block as prose, which the parser dropped. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 31 ++++++++++++------- .../test_agentcore_search_transformation.py | 31 +++++++++++++++++++ 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 94b69912faa..b53236347de 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -104,6 +104,13 @@ def _to_search_result(item: Mapping[str, object]) -> SearchResult: ) +def _result_items(parsed: object) -> tuple[Mapping[str, object], ...]: + items: Final = parsed.get("results", ()) if isinstance(parsed, Mapping) else parsed + if not isinstance(items, Sequence) or isinstance(items, (str, bytes)): + return () + return tuple(item for item in items if isinstance(item, Mapping)) + + def _parse_result_items(raw_text: object) -> tuple[Mapping[str, object], ...]: """ Parse one MCP text block into the search result objects it carries. @@ -117,10 +124,7 @@ def _parse_result_items(raw_text: object) -> tuple[Mapping[str, object], ...]: parsed: Final = json.loads(raw_text) except json.JSONDecodeError: return () - items: Final = parsed.get("results", ()) if isinstance(parsed, dict) else parsed - if not isinstance(items, Sequence) or isinstance(items, (str, bytes)): - return () - return tuple(item for item in items if isinstance(item, dict)) + return _result_items(parsed) def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: @@ -350,7 +354,9 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): The gateway returns JSON-RPC (as plain JSON or a single-message SSE stream) whose result.content[] text blocks contain a JSON list of - {title, url, date/publishedDate, text} entries. + {title, url, date/publishedDate, text} entries. Web-search connector + 1.1.0 and later repeat that list in result.structuredContent, which is + the only machine-readable copy when the text block holds prose instead. """ response_json: Final = self._parse_mcp_body(raw_response) @@ -370,14 +376,15 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", ) - return SearchResponse( - results=[ # mutable-ok: SearchResponse.results is a pydantic list field - _to_search_result(item) - for block in self._text_blocks(response_json) - for item in _parse_result_items(block.get("text")) - ], - object="search", + text_items: Final = tuple( + item for block in self._text_blocks(response_json) for item in _parse_result_items(block.get("text")) ) + structured: Final = result.get("structuredContent") if isinstance(result, Mapping) else None + items: Final = text_items or _result_items(structured) + + results: Final = [_to_search_result(item) for item in items] # mutable-ok: pydantic list field + + return SearchResponse(results=results, object="search") def _tool_error_message(self, response_json: Mapping[str, object]) -> str: texts: Final = tuple( diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 78808542c82..46f98356279 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -261,6 +261,37 @@ class TestAgentCoreSearch: with pytest.raises(Exception, match="AccessDeniedException"): config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + def test_transform_search_response_reads_structured_content(self): + """Connector 1.1.0+ puts the machine-readable results in structuredContent and may + leave the text block as prose, which must not come back as an empty result list.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [{"type": "text", "text": "Here is a prose summary of what I found."}], + "structuredContent": {"id": "824f89d0", "results": MCP_RESULTS}, + }, + } + ) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert [result.title for result in response.results] == ["Test Result 1", "Test Result 2"] + assert response.results[0].url == "https://example.com/1" + assert response.results[0].snippet == "Snippet for result 1" + assert response.results[0].date == "2026-06-16" + + def test_transform_search_response_does_not_duplicate_structured_content(self): + """1.1.0+ repeats the same results in both places, so parsing both would double them.""" + config = AgentCoreSearchConfig() + body = _mcp_response_body() + body["result"]["structuredContent"] = {"results": MCP_RESULTS} + mock_response = _make_mock_response(body) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + def test_transform_search_response_parses_crlf_framed_sse(self): """SSE streams may be CRLF framed; events must still split into separate events.""" config = AgentCoreSearchConfig() From 2a4598219df39f635bba411d2cd1bd27fc1fff04 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:43:22 -0700 Subject: [PATCH 42/88] feat(proxy): fast-fail validation for batch input files at /v1/files --- litellm/proxy/_types.py | 4 + .../batch_file_validation.py | 182 +++++++++++++++++ .../openai_files_endpoints/files_endpoints.py | 30 ++- litellm/proxy/proxy_server.py | 1 + .../test_files_batch_file_validation.py | 170 ++++++++++++++++ .../test_files_endpoint.py | 186 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 7 files changed, 567 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/openai_files_endpoints/batch_file_validation.py create mode 100644 tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8e57327b31b..6f352b73290 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2439,6 +2439,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="max request size in MB, if a request is larger than this size it will be rejected", ) + max_batch_file_size_mb: int | None = Field( + None, + description="max batch input file size in MB for /v1/files uploads with purpose=batch, if a file is larger than this size it will be rejected before being forwarded to the provider", + ) max_response_size_mb: int | None = Field( None, description="max response size in MB, if a response is larger than this size it will be rejected", diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py new file mode 100644 index 00000000000..bf6f6e2f829 --- /dev/null +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -0,0 +1,182 @@ +import json +from collections.abc import Iterator +from dataclasses import dataclass +from itertools import chain +from typing import BinaryIO, Final, NoReturn, assert_never + +from litellm.proxy._types import ProxyException + +BATCH_LINE_REQUIRED_KEYS: Final = ("custom_id", "method", "url", "body") +_MB: Final = 1024 * 1024 + + +@dataclass(frozen=True, slots=True) +class BatchFileTooLarge: + size_bytes: int + limit_mb: int + + +@dataclass(frozen=True, slots=True) +class BatchFileWrongExtension: + filename: str + + +@dataclass(frozen=True, slots=True) +class BatchFileEmpty: + pass + + +@dataclass(frozen=True, slots=True) +class BatchFileInvalidJsonLine: + line_number: int + + +@dataclass(frozen=True, slots=True) +class BatchFileLineNotObject: + line_number: int + + +@dataclass(frozen=True, slots=True) +class BatchFileMissingLineKey: + line_number: int + key: str + + +BatchFileValidationFailure = ( + BatchFileTooLarge + | BatchFileWrongExtension + | BatchFileEmpty + | BatchFileInvalidJsonLine + | BatchFileLineNotObject + | BatchFileMissingLineKey +) + + +def _file_size_bytes(file_source: bytes | BinaryIO) -> int: + if isinstance(file_source, bytes): + return len(file_source) + file_source.seek(0, 2) + size: Final = file_source.tell() + file_source.seek(0) + return size + + +def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: + if isinstance(file_source, bytes): + return iter(file_source.splitlines()) + file_source.seek(0) + return iter(file_source) + + +def _check_line(line_number: int, raw_line: bytes) -> BatchFileValidationFailure | None: + try: + parsed: Final = json.loads(raw_line) + except (json.JSONDecodeError, UnicodeDecodeError): + return BatchFileInvalidJsonLine(line_number=line_number) + if not isinstance(parsed, dict): + return BatchFileLineNotObject(line_number=line_number) + missing: Final = next((key for key in BATCH_LINE_REQUIRED_KEYS if key not in parsed), None) + if missing is None: + return None + return BatchFileMissingLineKey(line_number=line_number, key=missing) + + +def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | None: + content_lines: Final = ( + (line_number, raw_line) + for line_number, raw_line in enumerate(_iter_lines(file_source), start=1) + if raw_line.strip() + ) + first_line: Final = next(content_lines, None) + if first_line is None: + return BatchFileEmpty() + return next( + ( + failure + for line_number, raw_line in chain((first_line,), content_lines) + for failure in (_check_line(line_number, raw_line),) + if failure is not None + ), + None, + ) + + +def check_batch_file_upload( + filename: str | None, + file_source: bytes | BinaryIO, + max_batch_file_size_mb: int | None, +) -> BatchFileValidationFailure | None: + if filename is None or not filename.lower().endswith(".jsonl"): + return BatchFileWrongExtension(filename=filename or "") + if max_batch_file_size_mb is not None: + size_bytes: Final = _file_size_bytes(file_source) + if size_bytes > max_batch_file_size_mb * _MB: + return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) + scan_failure: Final = _scan_lines(file_source) + if not isinstance(file_source, bytes): + file_source.seek(0) + return scan_failure + + +def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> NoReturn: + match failure: + case BatchFileTooLarge(size_bytes=size_bytes, limit_mb=limit_mb): + raise ProxyException( + message=( + f"Batch input file is {size_bytes / _MB:.1f} MB, which exceeds the configured " + f"max_batch_file_size_mb of {limit_mb} MB. The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=413, + ) + case BatchFileWrongExtension(filename=filename): + raise ProxyException( + message=( + f"Invalid file format for Batch API: '{filename}'. " + "Batch input files must be .jsonl files. The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) + case BatchFileEmpty(): + raise ProxyException( + message="Batch input file has no request lines. The file was not forwarded to the provider.", + type="invalid_request_error", + param="file", + code=400, + ) + case BatchFileInvalidJsonLine(line_number=line_number): + raise ProxyException( + message=( + f"Batch input file line {line_number} is not valid JSON. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) + case BatchFileLineNotObject(line_number=line_number): + raise ProxyException( + message=( + f"Batch input file line {line_number} must be a JSON object. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) + case BatchFileMissingLineKey(line_number=line_number, key=key): + raise ProxyException( + message=( + f"Missing required parameter: '{key}' (batch input file line {line_number}). " + f"Each line must be a JSON object with keys {', '.join(BATCH_LINE_REQUIRED_KEYS)}. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param=key, + code=400, + ) + case _: + assert_never(failure) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 361b5b920e2..b7200de8fb6 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -21,6 +21,7 @@ from fastapi import ( UploadFile, status, ) +from pydantic import TypeAdapter import litellm from litellm import CreateFileRequest, get_secret_str @@ -41,6 +42,10 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.openai_files_endpoints.batch_file_validation import ( + check_batch_file_upload, + raise_batch_file_validation_failure, +) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, @@ -65,6 +70,8 @@ from litellm.types.llms.openai import ( router: Final = APIRouter() +_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None) + files_config = None @@ -361,18 +368,27 @@ async def create_file( # Prepare the data for forwarding - # Replace with: valid_purposes: Final = get_args(OpenAIFilesPurpose) if purpose not in valid_purposes: - raise HTTPException( - status_code=400, - detail={ - "error": f"Invalid purpose: {purpose}. Must be one of: {valid_purposes}", - }, + raise ProxyException( + message=f"Invalid purpose: {purpose}. Must be one of: {valid_purposes}", + type="invalid_request_error", + param="purpose", + code=400, ) # Cast purpose to OpenAIFilesPurpose type purpose = cast(OpenAIFilesPurpose, purpose) + if purpose == "batch": + batch_file_failure: Final = await asyncio.to_thread( + check_batch_file_upload, + file.filename, + file_source, + _MAX_BATCH_FILE_SIZE_MB_ADAPTER.validate_python(general_settings.get("max_batch_file_size_mb")), + ) + if batch_file_failure is not None: + raise_batch_file_validation_failure(batch_file_failure) + data = {} # Parse expires_after if provided @@ -552,6 +568,8 @@ async def create_file( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_file(): Exception occured - %s", e) + if isinstance(e, ProxyException): + raise e if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cae735be988..f08c9887a92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15689,6 +15689,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "max_parallel_requests": "Integer", "global_max_parallel_requests": "Integer", "max_request_size_mb": "Integer", + "max_batch_file_size_mb": "Integer", "max_response_size_mb": "Integer", "proxy_config_reload_interval_seconds": "Integer", "pass_through_endpoints": "PydanticModel", diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py new file mode 100644 index 00000000000..b4b0c5eb492 --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -0,0 +1,170 @@ +import io + +import pytest + +from litellm.proxy._types import ProxyException +from litellm.proxy.openai_files_endpoints.batch_file_validation import ( + BATCH_LINE_REQUIRED_KEYS, + BatchFileEmpty, + BatchFileInvalidJsonLine, + BatchFileLineNotObject, + BatchFileMissingLineKey, + BatchFileTooLarge, + BatchFileWrongExtension, + check_batch_file_upload, + raise_batch_file_validation_failure, +) + +VALID_LINE = ( + b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gpt-4.1-nano", "messages": [{"role": "user", "content": "hi"}]}}' +) + + +def test_valid_bytes_pass(): + assert check_batch_file_upload("batch.jsonl", VALID_LINE + b"\n" + VALID_LINE + b"\n", 10) is None + + +def test_valid_binaryio_passes_and_resets_position(): + handle = io.BytesIO(VALID_LINE + b"\n" + VALID_LINE + b"\n") + handle.seek(17) + assert check_batch_file_upload("batch.jsonl", handle, 10) is None + assert handle.tell() == 0 + + +def test_uppercase_extension_accepted(): + assert check_batch_file_upload("BATCH.JSONL", VALID_LINE, None) is None + + +@pytest.mark.parametrize("filename", ["batch.csv", "batch.json", "batch", None]) +def test_wrong_extension_rejected(filename): + assert check_batch_file_upload(filename, VALID_LINE, None) == BatchFileWrongExtension(filename=filename or "") + + +def test_size_over_cap_rejected_for_bytes(): + content = b"x" * (2 * 1024 * 1024) + assert check_batch_file_upload("batch.jsonl", content, 1) == BatchFileTooLarge( + size_bytes=len(content), limit_mb=1 + ) + + +def test_size_over_cap_rejected_for_binaryio(): + content = b"x" * (2 * 1024 * 1024) + assert check_batch_file_upload("batch.jsonl", io.BytesIO(content), 1) == BatchFileTooLarge( + size_bytes=len(content), limit_mb=1 + ) + + +def test_size_exactly_at_cap_allowed(): + line = VALID_LINE + b"\n" + padding_key = b'{"custom_id": "pad", "method": "POST", "url": "/v1/chat/completions", "body": {"note": "' + pad_line = padding_key + b"a" * (1024 * 1024 - len(line) - len(padding_key) - len(b'"}}\n')) + b'"}}\n' + content = line + pad_line + assert len(content) == 1024 * 1024 + assert check_batch_file_upload("batch.jsonl", content, 1) is None + + +def test_no_cap_skips_size_check(): + content = (VALID_LINE + b"\n") * 5000 + assert check_batch_file_upload("batch.jsonl", content, None) is None + + +@pytest.mark.parametrize("content", [b"", b"\n\n", b" \n\t\n"]) +def test_empty_file_rejected(content): + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileEmpty() + + +def test_invalid_json_line_rejected_with_line_number(): + content = VALID_LINE + b"\n" + b"not json at all\n" + VALID_LINE + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileInvalidJsonLine(line_number=2) + + +def test_non_utf8_line_rejected_as_invalid_json(): + assert check_batch_file_upload("batch.jsonl", b"\xff\xfe\x00\x01\n", None) == BatchFileInvalidJsonLine( + line_number=1 + ) + + +def test_non_object_line_rejected(): + content = VALID_LINE + b"\n" + b'["custom_id", "method"]\n' + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileLineNotObject(line_number=2) + + +@pytest.mark.parametrize("missing_key", BATCH_LINE_REQUIRED_KEYS) +def test_missing_required_key_rejected(missing_key): + import json + + line_dict = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4.1-nano"}, + } + del line_dict[missing_key] + content = VALID_LINE + b"\n" + json.dumps(line_dict).encode() + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileMissingLineKey( + line_number=2, key=missing_key + ) + + +def test_blank_lines_do_not_shift_line_numbers(): + content = b"\n" + VALID_LINE + b"\n\n" + b"broken\n" + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileInvalidJsonLine(line_number=4) + + +def test_failed_scan_leaves_handle_open_and_reset(): + handle = io.BytesIO(b"not json\n" + VALID_LINE + b"\n") + assert check_batch_file_upload("batch.jsonl", handle, None) == BatchFileInvalidJsonLine(line_number=1) + assert not handle.closed + assert handle.tell() == 0 + + +def test_scan_stops_at_first_failure(): + class ExplodingLines(io.BytesIO): + def __init__(self): + super().__init__(b"not json\n" + VALID_LINE + b"\n") + self.lines_read = 0 + + def __next__(self): + self.lines_read += 1 + return super().__next__() + + handle = ExplodingLines() + assert check_batch_file_upload("batch.jsonl", handle, None) == BatchFileInvalidJsonLine(line_number=1) + assert handle.lines_read == 1 + + +@pytest.mark.parametrize( + "failure, expected_code, expected_param, expected_fragments", + [ + ( + BatchFileTooLarge(size_bytes=220200960, limit_mb=10), + "413", + "file", + ("210.0 MB", "max_batch_file_size_mb", "10 MB", "not forwarded"), + ), + ( + BatchFileWrongExtension(filename="batch.csv"), + "400", + "file", + ("batch.csv", ".jsonl", "not forwarded"), + ), + (BatchFileEmpty(), "400", "file", ("no request lines", "not forwarded")), + (BatchFileInvalidJsonLine(line_number=3), "400", "file", ("line 3", "not valid JSON")), + (BatchFileLineNotObject(line_number=2), "400", "file", ("line 2", "JSON object")), + ( + BatchFileMissingLineKey(line_number=5, key="method"), + "400", + "method", + ("'method'", "line 5", "custom_id, method, url, body"), + ), + ], +) +def test_failures_map_to_openai_shaped_proxy_exceptions(failure, expected_code, expected_param, expected_fragments): + with pytest.raises(ProxyException) as exc_info: + raise_batch_file_validation_failure(failure) + assert exc_info.value.code == expected_code + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.param == expected_param + for fragment in expected_fragments: + assert fragment in exc_info.value.message diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index e363a266688..b02045a3713 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -225,7 +225,10 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router) assert response.status_code == 400 print(f"response: {response.json()}") - assert "Invalid purpose: my-bad-purpose" in response.json()["error"]["message"] + error = response.json()["error"] + assert "Invalid purpose: my-bad-purpose" in error["message"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "purpose" def test_get_file_content_rejects_raw_cloud_storage_uri(llm_router: Router): @@ -1599,7 +1602,7 @@ def _post_file_with_team_metadata( user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) app.dependency_overrides[user_api_key_auth] = lambda: user_key - test_file = ("mydata.jsonl", b'{"prompt": "Hello"}', "application/json") + test_file = ("mydata.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl") try: response = client.post( "/v1/files", @@ -1703,7 +1706,7 @@ def _post_file_raw( user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) app.dependency_overrides[user_api_key_auth] = lambda: user_key - test_file = ("mydata.jsonl", b'{"prompt": "Hello"}', "application/json") + test_file = ("mydata.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl") try: response = client.post( "/v1/files", @@ -2749,7 +2752,7 @@ def test_create_file_provider_only_resolves_named_vertex_credentials( try: response = client.post( "/v1/files", - files={"file": ("batch.jsonl", b"{}", "application/jsonl")}, + files={"file": ("batch.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl")}, data={"purpose": "batch"}, headers={ "Authorization": "Bearer test-key", @@ -2991,7 +2994,7 @@ def test_create_file_provider_only_skips_other_team_vertex_deployment( try: response = client.post( "/v1/files", - files={"file": ("batch.jsonl", b"{}", "application/jsonl")}, + files={"file": ("batch.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl")}, data={"purpose": "batch"}, headers={ "Authorization": "Bearer test-key", @@ -3341,3 +3344,176 @@ def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( assert response.status_code == 200, response.text mock_retrieve.assert_called_once() + + +VALID_BATCH_LINE = ( + b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}\n' +) + + +def _setup_batch_upload_endpoint(monkeypatch, llm_router: Router) -> list: + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + forwarded_calls: list = [] + + async def fake_route_create_file(**kwargs): + forwarded_calls.append(kwargs) + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + return forwarded_calls + + +def _teardown_batch_upload_endpoint(): + import litellm.proxy.proxy_server as ps + + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def test_create_file_batch_over_max_batch_file_size_mb_rejected_before_forwarding( + monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "max_batch_file_size_mb", 1) + + oversized = VALID_BATCH_LINE * (2 * 1024 * 1024 // len(VALID_BATCH_LINE) + 1) + try: + response = client.post( + "/v1/files", + files={"file": ("batch.jsonl", oversized, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 413, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert "max_batch_file_size_mb" in error["message"] + assert "1 MB" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_batch_under_max_batch_file_size_mb_forwards(monkeypatch, llm_router: Router): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "max_batch_file_size_mb", 1) + + try: + response = client.post( + "/v1/files", + files={"file": ("batch.jsonl", VALID_BATCH_LINE, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + assert len(forwarded_calls) == 1 + + +def test_create_file_batch_wrong_extension_rejected_before_forwarding(monkeypatch, llm_router: Router): + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + try: + response = client.post( + "/v1/files", + files={"file": ("batch.csv", VALID_BATCH_LINE, "text/csv")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert "batch.csv" in error["message"] + assert ".jsonl" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_batch_missing_line_key_rejected_before_forwarding(monkeypatch, llm_router: Router): + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + bad_line = b'{"custom_id": "req-1", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo"}}\n' + try: + response = client.post( + "/v1/files", + files={"file": ("batch.jsonl", VALID_BATCH_LINE + bad_line, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "method" + assert "line 2" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_batch_invalid_json_line_rejected_before_forwarding(monkeypatch, llm_router: Router): + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + try: + response = client.post( + "/v1/files", + files={"file": ("batch.jsonl", b"this is not jsonl\n", "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "file" + assert "line 1" in error["message"] + assert "not valid JSON" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_non_batch_purpose_skips_batch_validation(monkeypatch, llm_router: Router): + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + try: + response = client.post( + "/v1/files", + files={"file": ("notes.txt", b"plain text, not jsonl", "text/plain")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + assert len(forwarded_calls) == 1 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4c63586d4d6..d6569457cec 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24052,6 +24052,11 @@ export interface components { * @description require a key for all calls to proxy */ master_key?: string | null; + /** + * Max Batch File Size Mb + * @description max batch input file size in MB for /v1/files uploads with purpose=batch, if a file is larger than this size it will be rejected before being forwarded to the provider + */ + max_batch_file_size_mb?: number | null; /** * Max Parallel Requests * @description maximum parallel requests for each api key From 03ccfe9f985d842b8eff3166c0328a153f346463 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 21:45:20 +0000 Subject: [PATCH 43/88] docs(contributing): scope local unit test runs to the change Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- CONTRIBUTING.md | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d995ddcc87e..0821457584b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -13,8 +13,8 @@ Here are the core requirements for any PR submitted to LiteLLM: - [ ] **Add testing** - Adding at least 1 test is a hard requirement - [see details](#adding-testing) - [ ] **Ensure your PR passes all checks**: - - [ ] [Unit Tests](#running-unit-tests) - `make test-unit` - [ ] [Linting / Formatting](#running-linting-and-formatting-checks) - `make lint` + - [ ] [The tests covering your change](#running-unit-tests) pass, e.g. `uv run pytest tests/test_litellm/.py -v`. CI runs the full unit test matrix, so you don't need to run the whole suite locally #### UI PRs @@ -71,8 +71,8 @@ make format # Run all linting checks (matches CI exactly) make lint -# Run unit tests to ensure nothing is broken -make test-unit +# Run the tests covering your change (CI runs the full suite) +uv run pytest tests/test_litellm/.py -v # Commit your changes (must follow Conventional Commits — see above) git add . @@ -123,12 +123,13 @@ def test_your_feature(): ### Running Unit Tests -Run all unit tests (uses parallel execution for speed): - +Run the tests covering your change: ```bash -make test-unit +uv run pytest tests/test_litellm/test_your_file.py -v ``` +`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets, see `make help`), so reach for the group that covers your change instead of the whole suite. `make test-unit` runs everything when you really want it. + If you're running broader test suites, proxy tests, or anything that touches PostgreSQL-backed fixtures/plugins, install the full local test environment first: ```bash @@ -137,11 +138,6 @@ make install-test-deps This syncs the locked test environment used across the repo, including `psycopg` v3 plus `psycopg-binary` (used by `pytest-postgresql`), `psycopg2-binary` (used by some proxy E2E tests), and a generated Prisma client for DB-backed proxy tests, so pytest startup matches CI without manual package installs. -Run specific test files: -```bash -uv run pytest tests/test_litellm/test_your_file.py -v -``` - ### Running Linting and Formatting Checks Run all linting checks (matches CI exactly): From 81c975cff85877db5c41753bfd1d3eb43e2b8da8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:47:18 -0700 Subject: [PATCH 44/88] test(ui): cover the loading state while a newer user search is in flight --- .../user_search_modal.test.tsx | 21 ++++++++++++ .../create_key_button.integration.test.tsx | 32 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx index dfe6d762ac6..cf4461a2b7a 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -271,4 +271,25 @@ describe("UserSearchModal out-of-order search results", () => { expect(screen.queryByRole("option")).not.toBeInTheDocument(); expect(screen.getByText("No results")).toBeInTheDocument(); }); + + it("keeps loading while a newer search is still in flight", async () => { + const user = userEvent.setup(); + render(); + + const input = getEmailSearchInput(); + await user.click(input); + await user.type(input, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + + await user.type(input, "ce.smith@example.com"); + await waitFor(() => expect(answers.has("alice.smith@example.com")).toBe(true), { timeout: 3000 }); + + await answerFor("ali", []); + + expect(screen.getByText("Loading...")).toBeInTheDocument(); + expect(screen.queryByText("No results")).not.toBeInTheDocument(); + + await answerFor("alice.smith@example.com", [{ user_id: "u-smith", user_email: "alice.smith@example.com" }]); + await screen.findByRole("option", { name: "alice.smith@example.com" }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 8072a90fcad..dbeaf7998e6 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -863,6 +863,38 @@ describe("CreateKey", () => { expect(screen.queryByTitle("alice.jones@example.com (u-jones)")).not.toBeInTheDocument(); expect(screen.getByText("No users found")).toBeInTheDocument(); }); + + it("keeps searching while a newer search is still in flight", async () => { + const answers = new Map void>(); + vi.mocked(userFilterUICall).mockImplementation( + (_accessToken, params) => + new Promise((resolve) => { + answers.set(params.get("user_email") ?? "", resolve); + }) as never, + ); + + const user = userEvent.setup(); + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + const search = antdSearchInput(await screen.findByText("Type email to search for users")); + + await user.type(search, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + + await user.type(search, "ce.smith@example.com"); + await waitFor(() => expect(answers.has("alice.smith@example.com")).toBe(true), { timeout: 3000 }); + + await act(async () => { + answers.get("ali")?.([]); + }); + + expect(screen.getByText("Searching...")).toBeInTheDocument(); + expect(screen.queryByText("No users found")).not.toBeInTheDocument(); + + await act(async () => { + answers.get("alice.smith@example.com")?.([{ user_id: "u-smith", user_email: "alice.smith@example.com" }]); + }); + await screen.findByTitle("alice.smith@example.com (u-smith)"); + }); }); describe("created key display", () => { From 70a4f9a73aabd327b9c06e52bf5d47b437b76430 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:50:58 -0700 Subject: [PATCH 45/88] fix(search): refuse AgentCore credentials over plaintext HTTP A trusted hostname over plain http would expose the bearer token or a replayable SigV4 signature to network observers. Credentials now only ride https, with localhost exempt so local MCP stubs keep working. --- litellm/llms/bedrock/search/transformation.py | 15 +++++ .../test_agentcore_search_transformation.py | 60 +++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index b53236347de..920e566c9dd 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -87,6 +87,14 @@ def _gateway_host_match(api_base: str) -> re.Match[str] | None: return _GATEWAY_HOST_PATTERN.fullmatch(httpx.URL(api_base).host) +_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) + + +def _credential_safe_transport(api_base: str) -> bool: + url: Final = httpx.URL(api_base) + return url.scheme == "https" or url.host in _LOOPBACK_HOSTS + + def _string_field(item: Mapping[str, object], *keys: str) -> str | None: return next( (value for key in keys if isinstance(value := item.get(key), str) and value), @@ -259,6 +267,13 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): if not isinstance(request_data, dict): raise TypeError("AgentCore search expects a single dict request body") + if not _credential_safe_transport(api_base): + raise ValueError( + f"Refusing to send AgentCore credentials over plaintext HTTP to '{api_base}': a bearer " + "token or SigV4 signature would be readable in transit. Use an https gateway URL " + "(plain http is allowed only for localhost)." + ) + # Server-managed credentials only go to a trusted host, otherwise an # authenticated caller could point api_base at their own server (e.g. via # /search_tools/test_connection) and collect AGENTCORE_GATEWAY_TOKEN or a diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 46f98356279..950336c7ad0 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -403,6 +403,66 @@ class TestAgentCoreSearch: finally: os.environ.pop("AGENTCORE_GATEWAY_URL", None) + @pytest.mark.parametrize( + "plaintext_api_base", + [ + "http://gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", + "http://internal-gateway.corp/mcp", + ], + ) + def test_sign_request_refuses_server_token_over_plaintext_http(self, plaintext_api_base): + """A trusted hostname over plain http would expose the bearer token to + network observers, so credentials only ride https (or localhost).""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = plaintext_api_base + try: + with pytest.raises(ValueError, match="plaintext"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=plaintext_api_base, + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_refuses_sigv4_over_plaintext_http(self): + """Same for SigV4: a signature over plain http is replayable by observers.""" + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + with pytest.raises(ValueError, match="plaintext"): + config.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-1"}, + request_data={"jsonrpc": "2.0"}, + api_base="http://gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", + ) + mock_base_sign.assert_not_called() + + def test_sign_request_allows_plain_http_for_localhost(self): + """Local development against an MCP stub on 127.0.0.1 keeps working.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = "http://127.0.0.1:8931/mcp" + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="http://127.0.0.1:8931/mcp", + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + def test_sign_request_does_not_leak_bedrock_bearer_token(self): """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not replace SigV4 on requests to an AgentCore gateway.""" From 975e79dcefb599f7efd5cc6f1d83d77eff9f5959 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:55:53 -0700 Subject: [PATCH 46/88] fix(e2e): make concurrent replay consumption race-free Greptile flagged that lazy per-slug pool initialization could double-build under concurrent replay calls, splitting consumption across a discarded pool. Pools are now built once at ReplaySource construction and per-key consumption is a single atomic deque pop, with a barrier-synchronized regression test that fails 10/10 under the lazy-init mutant --- tests/e2e/fixture_transport.py | 28 +++++++++++-------- tests/e2e/test_fixture_transport.py | 43 +++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 11 deletions(-) diff --git a/tests/e2e/fixture_transport.py b/tests/e2e/fixture_transport.py index d49ac63787a..ce4eec701ca 100644 --- a/tests/e2e/fixture_transport.py +++ b/tests/e2e/fixture_transport.py @@ -391,18 +391,23 @@ def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle: @dataclass(slots=True) class ReplaySource: """One shared pool per test over a loaded bundle, so every client built in - the session consumes the same recorded interactions. Calls match by - canonical content key: order-independent across distinct keys (concurrent - tests interleave calls nondeterministically), FIFO within one key (a poll - loop replays its recorded responses in recorded order).""" + the session consumes the same recorded interactions. Every pool is built + once at construction and per-key consumption is a single atomic deque pop, + so concurrent replay calls never race. Calls match by canonical content + key: order-independent across distinct keys (concurrent tests interleave + calls nondeterministically), FIFO within one key (a poll loop replays its + recorded responses in recorded order).""" bundle: LoadedBundle - _pools: dict[str, dict[str, deque[Interaction]]] = field(default_factory=dict) + _pools: dict[str, dict[str, deque[Interaction]]] = field(init=False) + + def __post_init__(self) -> None: + self._pools = { + slug: _build_pool(recorded) for slug, recorded in self.bundle.interactions.items() + } def _pool(self, slug: str) -> dict[str, deque[Interaction]]: - if slug not in self._pools: - self._pools[slug] = _build_pool(self.bundle.interactions.get(slug, ())) - return self._pools[slug] + return self._pools.get(slug, {}) def next_interaction(self, request: RecordedRequest) -> Interaction: test_key: Final = current_test_key() @@ -412,12 +417,13 @@ class ReplaySource: queue: Final = pool.get(canonical.key) if queue is None: raise ReplayMiss(_miss_message(test_key, slug, canonical, self.bundle)) - if not queue: + try: + return queue.popleft() + except IndexError: raise ReplayMiss( f"replay exhausted for {test_key}: every recorded interaction for key " f"{canonical.key} is already consumed; re-record with E2E_FIXTURE_MODE=record" - ) - return queue.popleft() + ) from None def leftover_error(self, test_key: str) -> str | None: """Non-None when the test consumed fewer interactions than were recorded, diff --git a/tests/e2e/test_fixture_transport.py b/tests/e2e/test_fixture_transport.py index 814155f64d1..e61088d841c 100644 --- a/tests/e2e/test_fixture_transport.py +++ b/tests/e2e/test_fixture_transport.py @@ -14,6 +14,9 @@ pinned here too, including the stale message that names the bundle's age. from __future__ import annotations import hashlib +import sys +import threading +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path @@ -458,6 +461,46 @@ class TestReplayTransport: assert first == Success(status_code=200, data=Payload(value="first")) assert second == Success(status_code=200, data=Payload(value="second")) + def test_concurrent_replays_of_one_key_serve_each_recording_exactly_once(self, tmp_path: Path) -> None: + """A burst of parallel identical calls consumes one shared pool: no + response duplicated, none forgotten, nothing left over at teardown. + The tiny switch interval forces thread preemption inside pool setup + and consumption, so a non-atomic pool build or pop fails this test.""" + root = tmp_path / "bundle" + recorder = make_recorder(root) + for ordinal in range(32): + recorder.record( + test_key=current_test_key(), + request=recorded_request( + "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") + ), + response=RecordedResult(kind="success", status_code=200, data={"value": f"v{ordinal:02d}"}), + ) + source = replay_source(root) + replay: Transport = ReplayTransport(source=source, master_key="sk-1234") + barrier = threading.Barrier(8) + + def consume_one() -> str: + result = replay.get( + "/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload + ) + assert isinstance(result, Success) + return result.data.value + + def consume(_: int) -> tuple[str, ...]: + barrier.wait() + return tuple(consume_one() for _call in range(4)) + + previous_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + try: + with ThreadPoolExecutor(max_workers=8) as executor: + served = sorted(value for values in executor.map(consume, range(8)) for value in values) + finally: + sys.setswitchinterval(previous_interval) + assert served == [f"v{ordinal:02d}" for ordinal in range(32)] + assert source.leftover_error(current_test_key()) is None + class TestRecordedKeySets: def test_two_separate_recordings_of_one_flow_produce_identical_key_sets( From 710ef81a8050a7a9b1d2bd79e3d36b0e364e37bb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:57:19 -0700 Subject: [PATCH 47/88] fix(usage): keep responses usage SDK-parseable and complete streamed reasoning splits An unknown reasoning split now falls back to reasoning_tokens=0 in the chat-to-responses usage translation, since the OpenAI SDK requires output_tokens_details with an int reasoning_tokens, and the streaming chunk builder caps the tokenized reasoning estimate at completion_tokens and fills text_tokens with the remainder --- .../streaming_chunk_builder_utils.py | 7 +++- .../transformation.py | 25 +++++++------- .../test_streaming_chunk_builder_utils.py | 34 +++++++++++++++++++ .../test_litellm_completion_responses.py | 11 +++--- .../test_responses_api_bridge_non_stream.py | 29 +++++++++++++++- 5 files changed, 87 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 036ef3d5557..ee0518c4aec 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -987,7 +987,12 @@ class ChunkProcessor: returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - returned_usage.completion_tokens_details.reasoning_tokens = reasoning_tokens + capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) + returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens + if returned_usage.completion_tokens_details.text_tokens is None: + returned_usage.completion_tokens_details.text_tokens = ( + returned_usage.completion_tokens - capped_reasoning_tokens + ) if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b0099edc5dc..64084bfb063 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2303,18 +2303,19 @@ class LiteLLMCompletionResponsesConfig: # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: completion_details: Final = usage.completion_tokens_details - output_details_dict: Final[dict[str, int]] = {} - if hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None: - output_details_dict["reasoning_tokens"] = completion_details.reasoning_tokens - - if hasattr(completion_details, "text_tokens") and completion_details.text_tokens is not None: - output_details_dict["text_tokens"] = completion_details.text_tokens - - if hasattr(completion_details, "image_tokens") and completion_details.image_tokens is not None: - output_details_dict["image_tokens"] = completion_details.image_tokens - - if output_details_dict: - response_usage.output_tokens_details = OutputTokensDetails(**output_details_dict) + reasoning_token_count: Final = getattr(completion_details, "reasoning_tokens", None) + optional_output_details: Final[dict[str, int]] = { + field: value + for field, value in ( + ("text_tokens", getattr(completion_details, "text_tokens", None)), + ("image_tokens", getattr(completion_details, "image_tokens", None)), + ) + if value is not None + } + response_usage.output_tokens_details = OutputTokensDetails( + reasoning_tokens=reasoning_token_count if reasoning_token_count is not None else 0, + **optional_output_details, + ) return response_usage diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 222afdda3e7..0f21cce476b 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1308,3 +1308,37 @@ def test_count_reasoning_tokens_counts_visible_reasoning(): ) assert processor.count_reasoning_tokens(response) > 0 + + +@pytest.mark.parametrize( + "estimated_reasoning_tokens, expected_reasoning_tokens, expected_text_tokens", + [(40, 40, 60), (250, 100, 0)], +) +def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( + estimated_reasoning_tokens, expected_reasoning_tokens, expected_text_tokens +): + from litellm.types.utils import CompletionTokensDetailsWrapper + + chunk = ModelResponseStream( + id="chatcmpl-unknown-split", + model="claude-opus-4-8", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None, role=None))], + usage=Usage( + prompt_tokens=50, + completion_tokens=100, + total_tokens=150, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=None, text_tokens=None), + ), + ) + processor = ChunkProcessor(chunks=[chunk]) + + usage = processor.calculate_usage( + chunks=[chunk], + model="claude-opus-4-8", + completion_output="10", + reasoning_tokens=estimated_reasoning_tokens, + ) + + assert usage.completion_tokens == 100 + assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens + assert usage.completion_tokens_details.text_tokens == expected_text_tokens 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 0d4db2a0b11..aae053c2e8e 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 @@ -2638,10 +2638,10 @@ class TestUsageTransformation: assert response_usage.output_tokens_details.text_tokens == 50 assert response_usage.output_tokens_details.image_tokens == 100 - def test_reasoning_tokens_not_forced_to_zero_when_absent(self): - # Regression: previously the else branch wrote reasoning_tokens=0 even when - # completion_tokens_details had no reasoning (reasoning_tokens=None). That caused - # the proxy to always report reasoning_tokens=0 for non-thinking responses. + def test_reasoning_tokens_fall_back_to_zero_when_absent(self): + # The OpenAI SDK's ResponseUsage requires output_tokens_details.reasoning_tokens + # as an int, so an absent count degrades to 0 on the responses wire instead of + # dropping output_tokens_details and breaking SDK clients. usage = Usage( prompt_tokens=10, completion_tokens=50, @@ -2672,7 +2672,8 @@ class TestUsageTransformation: ) assert response_usage.output_tokens_details is not None - assert response_usage.output_tokens_details.reasoning_tokens is None + assert response_usage.output_tokens_details.reasoning_tokens == 0 + assert response_usage.output_tokens_details.text_tokens == 50 def test_reasoning_tokens_preserved_when_thinking_occurred(self): # Regression: reasoning_tokens must survive the chat->responses translation diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/test_litellm/test_responses_api_bridge_non_stream.py index c272b151865..08d55ee8290 100644 --- a/tests/test_litellm/test_responses_api_bridge_non_stream.py +++ b/tests/test_litellm/test_responses_api_bridge_non_stream.py @@ -297,7 +297,8 @@ def test_transform_usage_with_zero_values(): cached_tokens=0 is preserved (cache was available; nothing was cached). reasoning_tokens=0 is preserved the same way: an explicit provider-reported - zero passes through, while an absent value (None) is omitted. + zero passes through, while an absent value (None) falls back to 0 because the + Responses API wire contract requires reasoning_tokens as an int. """ completion_response = create_mock_completion_response( model="gpt-4", @@ -321,6 +322,32 @@ def test_transform_usage_with_zero_values(): print("✓ Transformation preserves explicit reasoning_tokens=0 and omits absent values") +def test_transform_usage_unknown_reasoning_split_keeps_output_tokens_details(): + """ + An unknown reasoning split (reasoning_tokens=None, text_tokens=None) must still + emit output_tokens_details with an integer reasoning_tokens: the OpenAI SDK's + ResponseUsage requires the field, so omitting it breaks /v1/responses clients. + """ + from openai.types.responses.response_usage import ( + OutputTokensDetails as OpenAISDKOutputTokensDetails, + ) + + from litellm.types.utils import CompletionTokensDetailsWrapper + + usage = Usage( + prompt_tokens=100, + completion_tokens=500, + total_tokens=600, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=None, text_tokens=None), + ) + + responses_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(usage) + + assert responses_usage.output_tokens_details is not None + assert responses_usage.output_tokens_details.reasoning_tokens == 0 + OpenAISDKOutputTokensDetails.model_validate(responses_usage.output_tokens_details.model_dump(exclude_none=True)) + + def test_input_tokens_details_requires_cached_tokens(): """ Test that InputTokensDetails has cached_tokens as an int with default value 0. From 680e4a5736d82f8b5774c362340b89e0e95a8b86 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:05:17 -0700 Subject: [PATCH 48/88] fix(files): treat nonpositive max_batch_file_size_mb as no cap --- .../proxy/openai_files_endpoints/batch_file_validation.py | 2 +- .../test_files_batch_file_validation.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index bf6f6e2f829..0aee5e8cc54 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -108,7 +108,7 @@ def check_batch_file_upload( ) -> BatchFileValidationFailure | None: if filename is None or not filename.lower().endswith(".jsonl"): return BatchFileWrongExtension(filename=filename or "") - if max_batch_file_size_mb is not None: + if max_batch_file_size_mb is not None and max_batch_file_size_mb > 0: size_bytes: Final = _file_size_bytes(file_source) if size_bytes > max_batch_file_size_mb * _MB: return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py index b4b0c5eb492..f5542fc0446 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -69,6 +69,12 @@ def test_no_cap_skips_size_check(): assert check_batch_file_upload("batch.jsonl", content, None) is None +@pytest.mark.parametrize("cap", [0, -3]) +def test_nonpositive_cap_disables_size_check(cap): + content = (VALID_LINE + b"\n") * 5000 + assert check_batch_file_upload("batch.jsonl", content, cap) is None + + @pytest.mark.parametrize("content", [b"", b"\n\n", b" \n\t\n"]) def test_empty_file_rejected(content): assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileEmpty() From 6b62b0b3861bbf5a69661f8703bd1b13b402a67e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:11:37 -0700 Subject: [PATCH 49/88] chore: make it more concise --- CONTRIBUTING.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0821457584b..9ef1d5ae2b8 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -128,7 +128,7 @@ Run the tests covering your change: uv run pytest tests/test_litellm/test_your_file.py -v ``` -`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets, see `make help`), so reach for the group that covers your change instead of the whole suite. `make test-unit` runs everything when you really want it. +`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets) on beefier boxes, so if, for whatever reason, you must run the whole suite, it's better to rely on CI to do that. If you're running broader test suites, proxy tests, or anything that touches PostgreSQL-backed fixtures/plugins, install the full local test environment first: From 103830ee86081046704d010cc84ae9e686dd40f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:13:56 -0700 Subject: [PATCH 50/88] fix(ui): drop the error toast for a superseded user search --- .../create_key_button.integration.test.tsx | 58 +++++++++++++++++++ .../organisms/create_key_button.tsx | 2 +- 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index dbeaf7998e6..d0add6fb12a 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { act, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import type { Team } from "../key_team_helpers/key_list"; import { keyCreateCall, keyCreateServiceAccountCall, modelAvailableCall, userFilterUICall } from "../networking"; +import { toast } from "@/lib/toast"; import CreateKey from "./create_key_button"; const state = vi.hoisted(() => ({ @@ -21,6 +22,16 @@ const state = vi.hoisted(() => ({ projects: [] as { project_id: string; project_alias: string; team_id?: string; models?: string[] }[], })); +vi.mock("@/lib/toast", () => ({ + toast: { + success: vi.fn(), + info: vi.fn(), + warning: vi.fn(), + error: vi.fn(), + fromError: vi.fn(), + dismiss: vi.fn(), + }, +})); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => state.authorized })); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: (capability: string) => state.can[capability] ?? true, @@ -225,6 +236,7 @@ describe("CreateKey", () => { .mockClear() .mockResolvedValue({ key: "sk-service-account", soft_budget: null }); vi.mocked(userFilterUICall).mockClear().mockResolvedValue([]); + vi.mocked(toast.fromError).mockClear(); vi.mocked(modelAvailableCall) .mockClear() .mockResolvedValue({ data: [{ id: "gpt-4" }] }); @@ -895,6 +907,52 @@ describe("CreateKey", () => { }); await screen.findByTitle("alice.smith@example.com (u-smith)"); }); + + it("only warns about a failed search when it is the one the box is waiting on", async () => { + const answers = new Map< + string, + { resolve: (users: { user_id: string; user_email: string }[]) => void; reject: (error: Error) => void } + >(); + vi.mocked(userFilterUICall).mockImplementation( + (_accessToken, params) => + new Promise((resolve, reject) => { + answers.set(params.get("user_email") ?? "", { resolve, reject }); + }) as never, + ); + + const user = userEvent.setup(); + renderCreateKey({ autoOpenCreate: true, prefillData: { owned_by: "another_user" } }); + const search = antdSearchInput(await screen.findByText("Type email to search for users")); + + await user.type(search, "ali"); + await waitFor(() => expect(answers.has("ali")).toBe(true), { timeout: 3000 }); + + await user.type(search, "ce.smith@example.com"); + await waitFor(() => expect(answers.has("alice.smith@example.com")).toBe(true), { timeout: 3000 }); + + await act(async () => { + answers + .get("alice.smith@example.com") + ?.resolve([{ user_id: "u-smith", user_email: "alice.smith@example.com" }]); + }); + await screen.findByTitle("alice.smith@example.com (u-smith)"); + + await act(async () => { + answers.get("ali")?.reject(new Error("search failed")); + }); + + expect(toast.fromError).not.toHaveBeenCalled(); + expect(screen.getByTitle("alice.smith@example.com (u-smith)")).toBeInTheDocument(); + + await user.type(search, "x"); + await waitFor(() => expect(answers.has("alice.smith@example.comx")).toBe(true), { timeout: 3000 }); + + await act(async () => { + answers.get("alice.smith@example.comx")?.reject(new Error("search failed")); + }); + + expect(toast.fromError).toHaveBeenCalledTimes(1); + }); }); describe("created key display", () => { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index b98cc4205bf..ea11a210470 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -587,7 +587,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setUserOptions(options); } catch (error) { console.error("Error fetching users:", error); - toast.fromError("Failed to search for users"); + if (isLatestSearch()) toast.fromError("Failed to search for users"); } finally { if (isLatestSearch()) setUserSearchLoading(false); } From 7744b9100ef68f43ab87ca8eb5b1b8d8b117ddc3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:14:20 -0700 Subject: [PATCH 51/88] docs(search): stop advertising yaml litellm_params knobs the search router drops --- .../agentcore_websearch_config.yaml | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml index bc29794e2a3..12402095c4d 100644 --- a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml +++ b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml @@ -18,16 +18,17 @@ search_tools: # The gateway exposes the connector as "___WebSearch". # Default is "web-search-tool___WebSearch", matching the target name used - # in the AWS docs' boto3/CLI setup examples. Set this ONLY if your target - # was created with a different name (misconfiguration surfaces as an MCP - # "tool not found" error): - # tool_name: MyWebSearchTarget___WebSearch + # in the AWS docs' boto3/CLI setup examples. If your target was created + # with a different name (misconfiguration surfaces as an MCP "tool not + # found" error), set the AGENTCORE_SEARCH_TOOL_NAME env var or pass + # tool_name in the request body. The search router forwards only + # search_provider / api_key / api_base from this litellm_params block, + # so a tool_name set here would be silently ignored. - # AWS_IAM gateway (default): SigV4-signed. Omit keys to use the standard - # AWS credential chain (env / profile / IRSA / instance role), or set them - # explicitly: - # aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID - # aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + # AWS_IAM gateway (default): SigV4-signed using the standard AWS + # credential chain (env / profile / IRSA / instance role). Explicit + # aws_access_key_id / aws_secret_access_key set here would be silently + # ignored for the same reason; pass them per request instead. # CUSTOM_JWT gateway alternative — OAuth2 bearer token instead of SigV4: # api_key: os.environ/AGENTCORE_GATEWAY_TOKEN From bd0c2fdb90b0ebb8365eca06256f578ceb82cc67 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 22:14:42 +0000 Subject: [PATCH 52/88] docs(pr-template): run only the tests covering your change, leave suites to CI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/pull_request_template.md | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 10266228b1f..cb5204b8e9f 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -53,6 +53,7 @@ After: the same request comes back with real token counts, so the dashboard show **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests +- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more. See [CONTRIBUTING.md](../CONTRIBUTING.md#running-unit-tests) - [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes) From c76223a3be7b20a418c08d547abe8f69d89a1594 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:19:43 -0700 Subject: [PATCH 53/88] chore: make it concise --- .github/pull_request_template.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index cb5204b8e9f..4e428d8cebf 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -53,8 +53,8 @@ After: the same request comes back with real token counts, so the dashboard show **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests -- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more. See [CONTRIBUTING.md](../CONTRIBUTING.md#running-unit-tests) -- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests) +- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more +- [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes) From 708ff0b91090eb7acd02e619bc8b20ebabf3f0f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:20:59 -0700 Subject: [PATCH 54/88] fix(proxy): retry end-user spend updates on Postgres deadlock instead of dropping them --- litellm/proxy/db/db_spend_update_writer.py | 6 ++ litellm/proxy/utils.py | 16 ++-- .../test_proxy_update_spend.py | 80 ++++++++++++++++++- 3 files changed, 90 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 19a82d556d8..65a271d4029 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1219,6 +1219,12 @@ class DBSpendUpdateWriter: is_retryable = isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) or PrismaDBExceptionHandler.is_deadlock_error(e) if not is_retryable or attempt >= n_retry_times: _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + verbose_proxy_logger.warning( + "Retrying spend update after retryable DB error (attempt %s/%s): %s", + attempt + 1, + n_retry_times, + e, + ) await asyncio.sleep(random.uniform(2**attempt, 2 ** (attempt + 1))) async def _commit_spend_updates_to_db( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a743526e975..41187af2bd8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -30,7 +30,6 @@ from litellm.constants import ( SPEND_LOG_WRITE_BATCH_MAX_BYTES, ) from litellm.proxy._types import ( - DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -5960,15 +5959,14 @@ class ProxyUpdateSpend: ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + await DBSpendUpdateWriter._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) @staticmethod async def update_spend_logs( diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index dd21bbc9e8a..7057a112c83 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -83,8 +83,8 @@ async def test_update_end_user_spend_retries_on_connect_error( mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: """``DB_RETRY_SAFE_ERROR_TYPES`` (ConnectError, statements provably never - sent) retries with backoff; once retries are exhausted the original - exception bubbles up via ``_raise_failed_update_spend_exception``. + sent) retries with jittered backoff; once retries are exhausted the + original exception bubbles up via ``_raise_failed_update_spend_exception``. """ import httpx import litellm.proxy.utils as utils_mod @@ -107,7 +107,8 @@ async def test_update_end_user_spend_retries_on_connect_error( proxy_logging_obj=proxy_logging, end_user_list_transactions={"u": 1.0}, ) - assert sleeps == [1.0] + assert len(sleeps) == 1 + assert 1.0 <= sleeps[0] <= 2.0 @pytest.mark.asyncio @@ -149,6 +150,79 @@ async def test_update_end_user_spend_non_connection_error_raises_immediately( ) +def _end_user_deadlock_error() -> Exception: + from prisma.errors import RawQueryError + + return RawQueryError(data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "LiteLLM_EndUserTable"}}}) + + +def _failing_tx(error: Exception) -> Any: + tx = MagicMock() + tx.__aenter__ = AsyncMock(side_effect=error) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +@pytest.mark.asyncio +async def test_update_end_user_spend_retries_on_deadlock_then_commits( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression for #27989: a Postgres deadlock (P2034/40P01) on the end-user + spend batch is retried with jittered backoff and the increments land, + instead of raising immediately and dropping the flushed spend.""" + sleeps: list[float] = [] + + async def _fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(asyncio, "sleep", _fake_sleep) + + batcher = MagicMock() + batcher.litellm_endusertable.upsert = MagicMock() + transaction = MagicMock() + transaction.batch_ = lambda: _AsyncCM(batcher) + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_end_user_deadlock_error()), _AsyncCM(transaction)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"end-user-1": 0.25}, + ) + + assert mock_prisma_client.db.tx.call_count == 2 + batcher.litellm_endusertable.upsert.assert_called_once() + assert batcher.litellm_endusertable.upsert.call_args.kwargs["where"] == {"user_id": "end-user-1"} + assert len(sleeps) == 1 + assert 1.0 <= sleeps[0] <= 2.0 + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_end_user_spend_raises_after_exhausting_deadlock_retries( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + from prisma.errors import RawQueryError + + monkeypatch.setattr(asyncio, "sleep", AsyncMock(return_value=None)) + mock_prisma_client.db.tx = MagicMock(side_effect=lambda timeout: _failing_tx(_end_user_deadlock_error())) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(RawQueryError): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=2, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"end-user-1": 0.25}, + ) + + assert mock_prisma_client.db.tx.call_count == 3 + + @pytest.mark.asyncio async def test_update_spend_logs_writes_batches_via_create_many( mock_prisma_client: Any, make_spend_log_row: Any From b39a339b7d9e67efcd8fddde8c289c59410d9160 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:21:06 -0700 Subject: [PATCH 55/88] fix(vertex_ai): apply regional endpoint uplift to cost tracking --- basedpyright-code-budget.json | 8 +- ci_cd/generate_model_prices_schema.py | 5 + litellm/cost_calculator.py | 15 ++ litellm/litellm_core_utils/litellm_logging.py | 28 +++- .../litellm_core_utils/llm_cost_calc/utils.py | 48 +++++++ litellm/llms/vertex_ai/cost_calculator.py | 19 ++- litellm/llms/vertex_ai/vertex_llm_base.py | 17 ++- ...odel_prices_and_context_window_backup.json | 25 ++++ litellm/proxy/spend_tracking/savings.py | 8 +- litellm/types/utils.py | 17 ++- litellm/utils.py | 1 + model_prices_and_context_window.json | 25 ++++ model_prices_and_context_window.schema.json | 5 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 132 ++++++++++++++++++ .../test_litellm_logging.py | 77 ++++++++++ .../proxy/spend_tracking/test_savings.py | 38 +++++ tests/test_litellm/test_cost_calculator.py | 69 +++++++++ type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 19 files changed, 516 insertions(+), 27 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 1ce71c5bd2c..32c77146b42 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5663 }, "reportMissingTypeArgument": { - "limit": 15557 + "limit": 15556 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39043 + "limit": 39030 }, "reportUnknownParameterType": { - "limit": 19887 + "limit": 19886 }, "reportUnknownVariableType": { - "limit": 30574 + "limit": 30573 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 153fbc0fdc2..252e3675329 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -145,6 +145,11 @@ NUMBER_KEYS: dict[str, JsonSchema] = { "minimum": 1, "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).", }, + "regional_endpoint_uplift_multiplier": { + "type": "number", + "minimum": 1, + "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).", + }, } COST_DESCRIPTIONS: dict[str, str] = { diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 7d7380665d3..8f7cd09d364 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -327,6 +327,8 @@ def cost_per_token( service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + ### VERTEX LOCATION ### + vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") response: Any | None = None, ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection @@ -587,6 +589,7 @@ def cost_per_token( prompt_characters=prompt_characters, completion_characters=completion_characters, usage=usage_block, + vertex_location=vertex_location, ) elif cost_router == "cost_per_token": return google_cost_per_token( @@ -594,6 +597,7 @@ def cost_per_token( custom_llm_provider=custom_llm_provider, usage=usage_block, service_tier=service_tier, + vertex_location=vertex_location, ) elif custom_llm_provider == "anthropic": return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) @@ -1071,6 +1075,7 @@ def _store_cost_breakdown_in_logging_obj( reasoning_cost: float | None = None, service_tier: str | None = None, data_residency: str | None = None, + vertex_location: str | None = None, ) -> None: """ Helper function to store cost breakdown in the logging object. @@ -1090,6 +1095,7 @@ def _store_cost_breakdown_in_logging_obj( margin_total_amount: Total margin added in USD service_tier: Tier the costs above were priced on, already resolved data_residency: Region uplift the costs above were priced on, already resolved + vertex_location: Vertex AI location the costs above were priced on, already resolved """ if litellm_logging_obj is None: return @@ -1113,6 +1119,7 @@ def _store_cost_breakdown_in_logging_obj( reasoning_cost=reasoning_cost, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, ) except Exception as breakdown_error: @@ -1149,6 +1156,8 @@ def completion_cost( service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + ### VERTEX LOCATION ### + vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -1577,6 +1586,7 @@ def completion_cost( rerank_billed_units=rerank_billed_units, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, ) @@ -1664,6 +1674,7 @@ def completion_cost( usage=cost_per_token_usage_object, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, ) _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost @@ -1686,6 +1697,7 @@ def completion_cost( reasoning_cost=_reasoning_cost, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, ) return _final_cost @@ -1765,6 +1777,8 @@ def response_cost_calculator( service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + ### VERTEX LOCATION ### + vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") ) -> float: """ Returns @@ -1797,6 +1811,7 @@ def response_cost_calculator( litellm_logging_obj=litellm_logging_obj, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, ) return response_cost except Exception as e: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 946110abf9e..e97abcf0af6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -13,7 +13,7 @@ import traceback from collections.abc import Callable, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache -from types import TracebackType +from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast from httpx import Response @@ -372,6 +372,24 @@ def _published_pricing(deployment_model: str | None) -> ModelInfo | None: return None +def _resolve_vertex_location_for_cost( + custom_llm_provider: str | None, + litellm_params: Mapping[str, object] | None, + model: str, +) -> str | None: + """ + The Vertex AI location a request was served from, resolved the same way + dispatch resolves it, so regional deployments price with the + regional-endpoint uplift. None for non-Vertex providers. + """ + if custom_llm_provider is None or not custom_llm_provider.startswith("vertex_ai"): + return None + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + configured_location: Final = VertexBase.safe_get_vertex_ai_location(litellm_params or MappingProxyType({})) + return VertexBase.get_vertex_region(configured_location, model) + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -1432,6 +1450,7 @@ class Logging(LiteLLMLoggingBaseClass): reasoning_cost: float | None = None, service_tier: str | None = None, data_residency: str | None = None, + vertex_location: str | None = None, ) -> None: """ Helper method to store cost breakdown in the logging object. @@ -1450,6 +1469,7 @@ class Logging(LiteLLMLoggingBaseClass): margin_total_amount: Total margin added in USD service_tier: Tier the costs above were priced on, already resolved data_residency: Region uplift the costs above were priced on, already resolved + vertex_location: Vertex AI location the costs above were priced on, already resolved """ self.cost_breakdown = CostBreakdown( @@ -1459,6 +1479,7 @@ class Logging(LiteLLMLoggingBaseClass): tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar, service_tier=service_tier, data_residency=data_residency, + vertex_location=vertex_location, ) if cache_read_cost is not None and cache_read_cost > 0: self.cost_breakdown["cache_read_cost"] = cache_read_cost @@ -1574,6 +1595,11 @@ class Logging(LiteLLMLoggingBaseClass): if hasattr(self, "litellm_params") and self.litellm_params else None ), + "vertex_location": _resolve_vertex_location_for_cost( + custom_llm_provider=self.model_call_details.get("custom_llm_provider", None), + litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None), + model=litellm_model_name or self.model, + ), } except Exception as e: # error creating kwargs for cost calculation debug_info = StandardLoggingModelCostFailureDebugInformation( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f73c4942a1c..dec35d16ea0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -757,6 +757,33 @@ def _get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | return 1.0 +def get_vertex_regional_endpoint_uplift(model_info: ModelInfo, vertex_location: str | None) -> float: + """ + Resolve the per-model uplift multiplier for Vertex AI non-global (regional and + multi-region) endpoints. + + Google prices every non-global endpoint at a flat premium over the global + endpoint (e.g. 1.10 = +10%) on all token types for the models that carry + regional pricing. The multiplier is stored on the model entry as + ``regional_endpoint_uplift_multiplier``. + + Returns 1.0 (no uplift) when ``vertex_location`` is ``None`` or ``"global"``, + or when the model has no multiplier configured. + """ + if vertex_location is None or vertex_location.lower() == "global": + return 1.0 + multiplier: Final = model_info.get("regional_endpoint_uplift_multiplier") + if multiplier is None: + return 1.0 + try: + return float(cast(float, multiplier)) + except (TypeError, ValueError): + verbose_logger.exception( + "Invalid regional_endpoint_uplift_multiplier for model; defaulting to 1.0", + ) + return 1.0 + + def get_provider_specific_geo_multiplier(model_info: ModelInfo, usage: Usage) -> float: """ Resolve the provider-specific regional pricing multiplier for the geo the @@ -798,6 +825,7 @@ def generic_cost_per_token( service_tier: str | None = None, data_residency: str | None = None, model_info: ModelInfo | None = None, + vertex_location: str | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -809,6 +837,9 @@ def generic_cost_per_token( - usage: LiteLLM Usage block, containing anthropic caching information - data_residency: optional OpenAI data-residency region (e.g. "eu", "us"), used to apply the per-model regional-processing uplift multiplier. + - vertex_location: optional Vertex AI location the request was served from + (e.g. "us-east5", "global"), used to apply the per-model + regional-endpoint uplift multiplier when non-global. Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -968,6 +999,14 @@ def generic_cost_per_token( prompt_cost *= uplift completion_cost *= uplift + ## VERTEX REGIONAL-ENDPOINT UPLIFT + # Applied as a flat multiplier across all token costs for the request + # when the Vertex AI endpoint serving it is non-global. + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + if vertex_uplift != 1.0: + prompt_cost *= vertex_uplift + completion_cost *= vertex_uplift + return prompt_cost, completion_cost @@ -988,6 +1027,7 @@ def get_token_type_cost_breakdown( usage: Usage, service_tier: str | None = None, data_residency: str | None = None, + vertex_location: str | None = None, ) -> TokenTypeCostBreakdown: """ Provider-agnostic cost of reasoning and cache tokens, derived from the usage @@ -1069,6 +1109,14 @@ def get_token_type_cost_breakdown( cache_read_cost *= uplift cache_creation_cost *= uplift + # Same flat uplift for Vertex AI non-global endpoints, keeping per-type + # costs reconciled with the totals for regional Vertex deployments. + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + if vertex_uplift != 1.0: + reasoning_cost *= vertex_uplift + cache_read_cost *= vertex_uplift + cache_creation_cost *= vertex_uplift + # Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals # apply, so cache and reasoning line items stay reconciled with them. geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index 86a5bb207ec..a9f5d77350c 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -7,6 +7,7 @@ from litellm import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import ( _is_above_128k, generic_cost_per_token, + get_vertex_regional_endpoint_uplift, ) from litellm.types.utils import ModelInfo, Usage @@ -63,6 +64,7 @@ def cost_per_character( usage: Usage, prompt_characters: float | None = None, completion_characters: float | None = None, + vertex_location: str | None = None, ) -> tuple[float, float]: """ Calculates the cost per character for a given VertexAI model, input messages, and response object. @@ -72,6 +74,8 @@ def cost_per_character( - custom_llm_provider: str, "vertex_ai-*" - prompt_characters: float, the number of input characters - completion_characters: float, the number of output characters + - vertex_location: the Vertex AI location serving the request; non-global + locations apply the model's regional-endpoint uplift multiplier Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -79,8 +83,6 @@ def cost_per_character( Raises: Exception if model requires >128k pricing, but model cost not mapped """ - model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - ## GET MODEL INFO model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) @@ -162,7 +164,10 @@ def cost_per_character( usage=usage, ) - return prompt_cost, completion_cost + # Applied once here; the cost_per_token fallbacks above are called without + # vertex_location so the uplift can never compound. + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + return prompt_cost * vertex_uplift, completion_cost * vertex_uplift def _handle_128k_pricing( @@ -196,6 +201,7 @@ def cost_per_token( custom_llm_provider: str, usage: Usage, service_tier: str | None = None, + vertex_location: str | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -207,6 +213,8 @@ def cost_per_token( - completion_tokens: float, the number of output tokens - service_tier: optional tier derived from Gemini trafficType ("priority" for ON_DEMAND_PRIORITY, "flex" for FLEX/batch). + - vertex_location: the Vertex AI location serving the request; non-global + locations apply the model's regional-endpoint uplift multiplier Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -222,14 +230,17 @@ def cost_per_token( input_cost_per_token_above_128k_tokens: Final = model_info.get("input_cost_per_token_above_128k_tokens") output_cost_per_token_above_128k_tokens: Final = model_info.get("output_cost_per_token_above_128k_tokens") if input_cost_per_token_above_128k_tokens is not None or output_cost_per_token_above_128k_tokens is not None: - return _handle_128k_pricing( + prompt_cost_128k, completion_cost_128k = _handle_128k_pricing( model_info=model_info, usage=usage, ) + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + return prompt_cost_128k * vertex_uplift, completion_cost_128k * vertex_uplift return generic_cost_per_token( model=model, custom_llm_provider=custom_llm_provider, usage=usage, service_tier=service_tier, + vertex_location=vertex_location, ) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 445e34966a9..0b3c003a60a 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -8,6 +8,7 @@ import asyncio import json import os import threading +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal from urllib.parse import urlparse @@ -68,7 +69,8 @@ class VertexBase: # re-acquire it without deadlocking the current thread. self._sync_refresh_lock = threading.RLock() - def get_vertex_region(self, vertex_region: str | None, model: str) -> str: + @staticmethod + def get_vertex_region(vertex_region: str | None, model: str) -> str: import litellm # Try to get supported_regions directly from model_cost @@ -1191,7 +1193,7 @@ class VertexBase: ) @staticmethod - def safe_get_vertex_ai_location(litellm_params: dict) -> str | None: + def safe_get_vertex_ai_location(litellm_params: Mapping[str, object]) -> str | None: """ Safely get Vertex AI location without mutating the litellm_params dict. @@ -1204,10 +1206,7 @@ class VertexBase: Returns: Vertex AI location/region or None """ - return ( - litellm_params.get("vertex_location") - or litellm_params.get("vertex_ai_location") - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - or get_secret_str("VERTEX_LOCATION") - ) + for configured in (litellm_params.get("vertex_location"), litellm_params.get("vertex_ai_location")): + if isinstance(configured, str) and configured: + return configured + return litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") or get_secret_str("VERTEX_LOCATION") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f9027313b..52f1763a5ba 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19624,6 +19624,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -19679,6 +19680,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -19733,6 +19735,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -38685,6 +38688,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -38708,6 +38712,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -38923,6 +38928,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38951,6 +38957,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38970,6 +38977,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39000,6 +39008,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39030,6 +39039,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39061,6 +39071,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39092,6 +39103,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -39123,6 +39135,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -39154,6 +39167,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39186,6 +39200,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39218,6 +39233,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39250,6 +39266,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39298,6 +39315,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -39310,6 +39328,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39342,6 +39361,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-6": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -39388,6 +39408,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -39794,6 +39815,7 @@ "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, "output_cost_per_token_priority": 2.7e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -39849,6 +39871,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -47014,6 +47037,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -47046,6 +47070,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-6@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 448723ab3bc..997180efdde 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -130,6 +130,7 @@ class PricingBasis(NamedTuple): service_tier: str | None = None data_residency: str | None = None + vertex_location: str | None = None _STANDARD_RATES: Final = PricingBasis() @@ -141,8 +142,8 @@ def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis: Rows written before this field shipped carry neither key, and there is no backfill: they price at standard rates, which is what they already did. - Both values survive a JSON round trip on the way here, so neither is guaranteed to be - a string. `generic_cost_per_token` calls `.lower()` on both without a type check, and + These values survive a JSON round trip on the way here, so none is guaranteed to be + a string. `generic_cost_per_token` calls `.lower()` on them without a type check, and the resulting `AttributeError` would be swallowed into a silent zero by the caller's `except`, so anything that is not a string is dropped here instead. """ @@ -150,9 +151,11 @@ def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis: return _STANDARD_RATES service_tier: Final = cost_breakdown.get("service_tier") data_residency: Final = cost_breakdown.get("data_residency") + vertex_location: Final = cost_breakdown.get("vertex_location") return PricingBasis( service_tier=service_tier if isinstance(service_tier, str) else None, data_residency=data_residency if isinstance(data_residency, str) else None, + vertex_location=vertex_location if isinstance(vertex_location, str) else None, ) @@ -193,6 +196,7 @@ def _cost_of_usage( service_tier=basis.service_tier, data_residency=basis.data_residency, model_info=model_info, + vertex_location=basis.vertex_location, ) except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings verbose_proxy_logger.debug( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 07005d7f9ad..5f629eb129f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -248,6 +248,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): regional_processing_uplift_multiplier_us: ( float | None ) # OpenAI US data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) + regional_endpoint_uplift_multiplier: ReadOnly[ + float | None + ] # Vertex AI non-global (regional) endpoint uplift multiplier applied to all token costs (e.g. 1.10 = +10%) output_cost_per_character: float | None # only for vertex ai models output_cost_per_audio_token: float | None output_cost_per_token_above_128k_tokens: float | None # only for vertex ai models @@ -3113,16 +3116,17 @@ class CostBreakdown(TypedDict, total=False): """ Detailed cost breakdown for a request. - ``service_tier`` and ``data_residency`` record the pricing basis the cost was - computed on, not what the caller asked for. A consumer that has to price a - counterfactual against this request (what another model would have charged for - it) needs the same basis to compare like with like, and re-deriving it from the - request is not possible after the fact: the tier the biller used comes from - ``optional_params``, which no log record carries. + ``service_tier``, ``data_residency``, and ``vertex_location`` record the pricing + basis the cost was computed on, not what the caller asked for. A consumer that has + to price a counterfactual against this request (what another model would have + charged for it) needs the same basis to compare like with like, and re-deriving it + from the request is not possible after the fact: the tier the biller used comes + from ``optional_params``, which no log record carries. """ service_tier: str | None data_residency: str | None + vertex_location: ReadOnly[str | None] input_cost: float # Cost of raw (non-cached) input tokens only cache_read_cost: float # Cost of cache-read tokens (discounted rate) cache_creation_cost: float # Cost of cache-write tokens (premium rate) @@ -3388,6 +3392,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): annotation_cost_per_page: float | None = None regional_processing_uplift_multiplier_eu: float | None = None regional_processing_uplift_multiplier_us: float | None = None + regional_endpoint_uplift_multiplier: float | None = None @classmethod def strip_custom_pricing_fields(cls, model_info: dict[str, Any]) -> dict[str, Any]: diff --git a/litellm/utils.py b/litellm/utils.py index a7b70c4129a..8be2c98fc68 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5662,6 +5662,7 @@ def _get_model_info_helper( regional_processing_uplift_multiplier_us=_model_info.get( "regional_processing_uplift_multiplier_us", None ), + regional_endpoint_uplift_multiplier=_model_info.get("regional_endpoint_uplift_multiplier", None), output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None), output_cost_per_character=_model_info.get("output_cost_per_character", None), output_cost_per_reasoning_token=_model_info.get("output_cost_per_reasoning_token", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 07f9027313b..52f1763a5ba 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19624,6 +19624,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -19679,6 +19680,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -19733,6 +19735,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -38685,6 +38688,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -38708,6 +38712,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/haiku-4-5", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -38923,6 +38928,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38951,6 +38957,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, + "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38970,6 +38977,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39000,6 +39008,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39030,6 +39039,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39061,6 +39071,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -39092,6 +39103,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -39123,6 +39135,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -39154,6 +39167,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39186,6 +39200,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39218,6 +39233,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39250,6 +39266,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39298,6 +39315,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -39310,6 +39328,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39342,6 +39361,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-6": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -39388,6 +39408,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -39794,6 +39815,7 @@ "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, "output_cost_per_token_priority": 2.7e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -39849,6 +39871,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "regional_endpoint_uplift_multiplier": 1.1, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -47014,6 +47037,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -47046,6 +47070,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-6@default": { + "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index cd02fde595f..82854a3b717 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -514,6 +514,11 @@ "type": "object", "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema)." }, + "regional_endpoint_uplift_multiplier": { + "type": "number", + "minimum": 1, + "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%)." + }, "regional_processing_uplift_multiplier_eu": { "type": "number", "minimum": 1, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 1826f56d667..06be96fefdf 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2324,6 +2324,87 @@ def test_data_residency_composes_with_service_tier(_local_model_cost_map): assert priority_eu_total == pytest.approx(priority_base_total * 1.10, rel=1e-9) +@pytest.mark.parametrize("model", ["gemini-3.5-flash", "claude-haiku-4-5@20251001"]) +@pytest.mark.parametrize("vertex_location", ["us-central1", "us-east5", "europe-west1", "asia-southeast1"]) +def test_vertex_regional_location_applies_uplift(vertex_location, model, _local_model_cost_map): + """Google bills every non-global Vertex endpoint at 1.1x the global rate for GA + Gemini 3+ and regional-pricing Claude models, so a request served from a regional + location must cost 1.1x what the same usage costs on the global endpoint.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="vertex_ai") + regional = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="vertex_ai", + vertex_location=vertex_location, + ) + + base_total = base[0] + base[1] + regional_total = regional[0] + regional[1] + + assert base_total > 0 + assert regional_total == pytest.approx(base_total * 1.10, rel=1e-9) + assert regional[0] == pytest.approx(base[0] * 1.10, rel=1e-9) + assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9) + + +@pytest.mark.parametrize("vertex_location", [None, "global", "GLOBAL"]) +def test_vertex_global_or_absent_location_no_uplift(vertex_location, _local_model_cost_map): + """The global endpoint prices at the base rate, whatever the casing, and an + unresolved location must never uplift.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token( + model="claude-haiku-4-5@20251001", usage=usage, custom_llm_provider="vertex_ai" + ) + located = generic_cost_per_token( + model="claude-haiku-4-5@20251001", + usage=usage, + custom_llm_provider="vertex_ai", + vertex_location=vertex_location, + ) + + assert base == located + + +@pytest.mark.parametrize("model", ["claude-opus-4-1", "gemini-2.0-flash-001"]) +def test_vertex_location_no_uplift_for_uniformly_priced_model(model, _local_model_cost_map): + """Models Google prices uniformly across endpoints (Gemini 2.x, Claude Opus 4.1 + and older) carry no multiplier and must not move with the location.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="vertex_ai") + regional = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="vertex_ai", + vertex_location="us-east5", + ) + + assert base == regional, f"{model} should not have a regional-endpoint uplift" + + +def test_vertex_uplift_invalid_multiplier_defaults_to_one(): + """A malformed multiplier in the cost map degrades to base pricing, never raises.""" + from litellm.litellm_core_utils.llm_cost_calc.utils import ( + get_vertex_regional_endpoint_uplift, + ) + + assert ( + get_vertex_regional_endpoint_uplift( + {"regional_endpoint_uplift_multiplier": "not-a-number"}, "us-east5" + ) + == 1.0 + ) + + def test_priority_service_tier_above_threshold_uses_priority_tier_rates_for_cached_tokens( _local_model_cost_map, ): @@ -2877,6 +2958,57 @@ def test_token_type_cost_breakdown_applies_regional_uplift(): assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost) +def test_token_type_cost_breakdown_applies_vertex_regional_uplift(): + """ + Non-global Vertex endpoints apply a flat 1.1x uplift to every token cost. The + per-type breakdown must apply the same uplift via vertex_location so it stays + reconciled with the uplifted input_cost/output_cost totals, instead of being + logged at the global rate. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "claude-haiku-4-5@20251001" + custom_llm_provider = "vertex_ai" + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=400, text_tokens=600 + ), + ) + + model_info = litellm.get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + uplift = model_info["regional_endpoint_uplift_multiplier"] + assert uplift > 1.0 + + base = get_token_type_cost_breakdown( + model=model, custom_llm_provider=custom_llm_provider, usage=usage + ) + regional = get_token_type_cost_breakdown( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + vertex_location="us-east5", + ) + + assert base.cache_read_cost > 0 + assert regional.cache_read_cost == pytest.approx(base.cache_read_cost * uplift) + + # The uplifted breakdown must still reconcile with the uplifted totals. + prompt_cost, _completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + vertex_location="us-east5", + ) + text_input_cost = 600 * model_info["input_cost_per_token"] * uplift + assert text_input_cost + regional.cache_read_cost == pytest.approx(prompt_cost) + + def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch): """ Anthropic's regional (geo) uplift lives in provider_specific_entry and is diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 54016470f8b..01e5ec26b1e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4958,3 +4958,80 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj): raw_api_base = logging_obj.model_call_details["raw_request_typed_dict"]["raw_request_api_base"] assert _GEMINI_KEY not in raw_api_base assert "key=*****" in raw_api_base + + +def test_resolve_vertex_location_for_cost(): + """Vertex requests resolve the serving location the way dispatch does; other providers get None.""" + from litellm.litellm_core_utils.litellm_logging import ( + _resolve_vertex_location_for_cost, + ) + + assert _resolve_vertex_location_for_cost("openai", {"vertex_location": "us-east5"}, "gpt-4o") is None + assert _resolve_vertex_location_for_cost(None, {}, "gemini-3.5-flash") is None + assert ( + _resolve_vertex_location_for_cost("vertex_ai", {"vertex_location": "us-east5"}, "gemini-3.5-flash") + == "us-east5" + ) + assert ( + _resolve_vertex_location_for_cost("vertex_ai", {"vertex_location": "global"}, "gemini-3.5-flash") == "global" + ) + assert ( + _resolve_vertex_location_for_cost( + "vertex_ai_beta", {"vertex_ai_location": "europe-west1"}, "claude-haiku-4-5@20251001" + ) + == "europe-west1" + ) + + +def test_resolve_vertex_location_for_cost_default_region(monkeypatch): + """With no location configured anywhere, resolution lands on the dispatch default us-central1.""" + from litellm.litellm_core_utils.litellm_logging import ( + _resolve_vertex_location_for_cost, + ) + + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) + monkeypatch.setattr(litellm, "vertex_location", None) + + assert _resolve_vertex_location_for_cost("vertex_ai", {}, "gemini-3.5-flash") == "us-central1" + assert _resolve_vertex_location_for_cost("vertex_ai", None, "gemini-3.5-flash") == "us-central1" + + +def test_set_cost_breakdown_stores_vertex_location(): + """vertex_location is recorded in the pricing basis, None for non-vertex requests.""" + from datetime import datetime + + logging_obj = LitellmLogging( + model="vertex_ai/claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="vertex-location-set", + function_id="f", + ) + logging_obj.set_cost_breakdown( + input_cost=0.001, + output_cost=0.002, + total_cost=0.003, + cost_for_built_in_tools_cost_usd_dollar=0.0, + vertex_location="us-east5", + ) + assert logging_obj.cost_breakdown["vertex_location"] == "us-east5" + + no_location = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="vertex-location-absent", + function_id="f", + ) + no_location.set_cost_breakdown( + input_cost=0.001, + output_cost=0.002, + total_cost=0.003, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) + assert no_location.cost_breakdown.get("vertex_location") is None diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index 1435547c434..9006288bdae 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -841,6 +841,44 @@ def test_the_baseline_is_priced_on_the_basis_the_request_was_billed_at(basis, ex assert reported == pytest.approx(expected_multiplier * baseline - served) +def test_the_baseline_is_priced_on_the_vertex_location_the_request_was_billed_at(monkeypatch): + """A request served from a regional Vertex endpoint was billed with the + regional-endpoint uplift, so the counterfactual single-model operator would + have paid it too. The served model carries no uplift field, so only the + baseline moves with the recorded location.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + gemini = litellm.get_model_info("gemini-3.5-flash", "vertex_ai") + haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") + assert gemini.get("regional_endpoint_uplift_multiplier") == 1.1 + assert haiku.get("regional_endpoint_uplift_multiplier") is None, "served model must not move with the basis" + + usage = _usage(fresh=20_000, cached=0, written=0, out=1_000) + served = 20_000 * haiku["input_cost_per_token"] + 1_000 * haiku["output_cost_per_token"] + baseline = 20_000 * gemini["input_cost_per_token"] + 1_000 * gemini["output_cost_per_token"] + + regional = compute_autorouter_savings( + baseline_model="vertex_ai/gemini-3.5-flash", + selected_model="claude-haiku-4-5", + selected_provider="anthropic", + usage=usage, + conversation_continuing=False, + cost_breakdown=_breakdown(served, vertex_location="us-east5"), + ) + global_endpoint = compute_autorouter_savings( + baseline_model="vertex_ai/gemini-3.5-flash", + selected_model="claude-haiku-4-5", + selected_provider="anthropic", + usage=usage, + conversation_continuing=False, + cost_breakdown=_breakdown(served, vertex_location="global"), + ) + + assert regional == pytest.approx(1.1 * baseline - served) + assert global_endpoint == pytest.approx(baseline - served) + + def test_a_baseline_recorded_on_the_decision_turns_the_driver_on(): """An operator who configures nothing still sees the driver work.""" result = compute_savings_spend( diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 75c90d793fe..6deaf5479e0 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1742,6 +1742,75 @@ def test_azure_ai_cache_cost_calculation(): ), f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" +def test_vertex_regional_deployment_costs_uplift_over_global(monkeypatch): + """ + Regression for https://github.com/BerriAI/litellm/issues/34393: two Vertex + deployments differing only in vertex_location must not price identically. + Google bills non-global endpoints at 1.1x for regional-pricing models, so the + regional request costs 1.1x the global one for the exact same usage, through + both vertex cost routes (Claude via cost_per_token, Gemini via + cost_per_character's token fallback). + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + usage = Usage(prompt_tokens=15, completion_tokens=5, total_tokens=20) + for model in ("claude-haiku-4-5@20251001", "gemini-3.5-flash"): + global_prompt, global_completion = cost_per_token( + model=model, + custom_llm_provider="vertex_ai", + usage_object=usage, + vertex_location="global", + ) + regional_prompt, regional_completion = cost_per_token( + model=model, + custom_llm_provider="vertex_ai", + usage_object=usage, + vertex_location="us-east5", + ) + global_total = global_prompt + global_completion + regional_total = regional_prompt + regional_completion + assert global_total > 0 + assert regional_total == pytest.approx(global_total * 1.10, rel=1e-9), ( + f"{model}: regional Vertex request must cost 1.1x the global one" + ) + + +def test_vertex_uplift_composes_with_above_128k_pricing(monkeypatch): + """The regional-endpoint uplift multiplies whatever rate the request priced at, + including the above-128k dynamic rates, so a synthetic model carrying both keys + prices regional above-128k usage at 1.1x the above-128k rate.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.model_cost["vertex_ai/fake-regional-128k-model"] = { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "input_cost_per_token_above_128k_tokens": 2e-06, + "output_cost_per_token_above_128k_tokens": 4e-06, + "regional_endpoint_uplift_multiplier": 1.1, + } + + usage = Usage(prompt_tokens=200_000, completion_tokens=10, total_tokens=200_010) + global_prompt, global_completion = cost_per_token( + model="fake-regional-128k-model", + custom_llm_provider="vertex_ai", + usage_object=usage, + vertex_location="global", + ) + regional_prompt, regional_completion = cost_per_token( + model="fake-regional-128k-model", + custom_llm_provider="vertex_ai", + usage_object=usage, + vertex_location="europe-west1", + ) + + assert global_prompt == pytest.approx(200_000 * 2e-06, rel=1e-9) + assert regional_prompt == pytest.approx(global_prompt * 1.10, rel=1e-9) + assert regional_completion == pytest.approx(global_completion * 1.10, rel=1e-9) + + def test_cost_discount_vertex_ai(): """ Test that cost discount is applied correctly for Vertex AI provider diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6f77a621a9e..56ffbe9fde8 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22809 + "limit": 22808 }, "LIT002": { "limit": 26878 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5794c60e97b..3cc7397f5bb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27559,6 +27559,8 @@ export interface components { quality_router_default_model?: string | null; /** Region Name */ region_name?: string | null; + /** Regional Endpoint Uplift Multiplier */ + regional_endpoint_uplift_multiplier?: number | null; /** Regional Processing Uplift Multiplier Eu */ regional_processing_uplift_multiplier_eu?: number | null; /** Regional Processing Uplift Multiplier Us */ @@ -36689,6 +36691,8 @@ export interface components { quality_router_default_model?: string | null; /** Region Name */ region_name?: string | null; + /** Regional Endpoint Uplift Multiplier */ + regional_endpoint_uplift_multiplier?: number | null; /** Regional Processing Uplift Multiplier Eu */ regional_processing_uplift_multiplier_eu?: number | null; /** Regional Processing Uplift Multiplier Us */ From 0f4c2d71fb04681fef221231d88695a6b8c01e8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:22:43 -0700 Subject: [PATCH 56/88] test(files): hoist shared batch line fixture into one constant --- .../test_files_endpoint.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index b02045a3713..99fb19f0d60 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -31,6 +31,11 @@ from litellm.caching.caching import DualCache from litellm.proxy.proxy_server import hash_token from litellm.proxy.utils import ProxyLogging +VALID_BATCH_LINE = ( + b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}\n' +) + @pytest.fixture def llm_router() -> Router: @@ -1602,7 +1607,7 @@ def _post_file_with_team_metadata( user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) app.dependency_overrides[user_api_key_auth] = lambda: user_key - test_file = ("mydata.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl") + test_file = ("mydata.jsonl", VALID_BATCH_LINE, "application/jsonl") try: response = client.post( "/v1/files", @@ -1706,7 +1711,7 @@ def _post_file_raw( user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) app.dependency_overrides[user_api_key_auth] = lambda: user_key - test_file = ("mydata.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl") + test_file = ("mydata.jsonl", VALID_BATCH_LINE, "application/jsonl") try: response = client.post( "/v1/files", @@ -2752,7 +2757,7 @@ def test_create_file_provider_only_resolves_named_vertex_credentials( try: response = client.post( "/v1/files", - files={"file": ("batch.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl")}, + files={"file": ("batch.jsonl", VALID_BATCH_LINE, "application/jsonl")}, data={"purpose": "batch"}, headers={ "Authorization": "Bearer test-key", @@ -2994,7 +2999,7 @@ def test_create_file_provider_only_skips_other_team_vertex_deployment( try: response = client.post( "/v1/files", - files={"file": ("batch.jsonl", b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}', "application/jsonl")}, + files={"file": ("batch.jsonl", VALID_BATCH_LINE, "application/jsonl")}, data={"purpose": "batch"}, headers={ "Authorization": "Bearer test-key", @@ -3346,12 +3351,6 @@ def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( mock_retrieve.assert_called_once() -VALID_BATCH_LINE = ( - b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions",' - b' "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}}\n' -) - - def _setup_batch_upload_endpoint(monkeypatch, llm_router: Router) -> list: import litellm.proxy.proxy_server as ps from litellm.proxy._types import LitellmUserRoles From 6ca48efc8bdc103475271838b6c7174782276822 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 19 Aug 2026 15:32:14 -0700 Subject: [PATCH 57/88] feat(cli): add `lite login --config-claude` to wire Claude Code at login (#37507) `lite up` already patches ~/.claude/settings.json, but only for as long as it runs in the foreground, and it restores the original file on exit. Users proxying Claude Code through LiteLLM therefore have to re-wire it by hand after every login. --config-claude makes that write persistent. It reuses the settings shape `lite up` writes (env.ANTHROPIC_BASE_URL plus an apiKeyHelper invocation), preserves every unrelated key, creates the file when missing, and writes it atomically with owner-only permissions. Plain `lite login` is unchanged. Reaching the credential through apiKeyHelper rather than copying it into the file means a later login refreshes it with no further action, and keeps the short-lived CLI token out of settings.json entirely. The shared parts of the settings-file handling move from up.py into a new claude_settings.py, since up.py imports auth.py and so auth.py cannot import up.py back. That module now also owns the registry of commands that can be temporarily managing the file, so the persistent write refuses while either `lite up` or `lite autoroute up` holds a backup it would later restore over this write. Because this write has no backup and no `lite down`, it is stricter than `lite up` about the user's file: it writes through a symlinked settings.json rather than replacing the link with a regular file, and it refuses rather than silently discarding a non-object `env` value. Also fixes the apiKeyHelper command itself: --base-url belongs to the top-level `lite` group, so `lite auth print-token --base-url X` is rejected by click with "No such option". Every settings file `lite up` has written carries that malformed command, which makes the helper return nothing and every Claude Code request lose its token. The existing tests only string-matched the generated command, so the new tests parse it through the real CLI instead. --- litellm/proxy/client/cli/README.md | 14 + litellm/proxy/client/cli/commands/auth.py | 34 ++- .../client/cli/commands/autoroute/commands.py | 21 +- .../client/cli/commands/claude_settings.py | 155 ++++++++++ litellm/proxy/client/cli/commands/up.py | 80 +----- .../proxy/client/cli/test_auth_commands.py | 83 +++++- .../proxy/client/cli/test_claude_settings.py | 269 ++++++++++++++++++ .../proxy/client/cli/test_up_commands.py | 59 +++- 8 files changed, 633 insertions(+), 82 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/claude_settings.py create mode 100644 tests/test_litellm/proxy/client/cli/test_claude_settings.py diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 72ed67728d8..9fe52b7a27d 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -521,6 +521,20 @@ This is a one-time file patch and restore, not a live traffic interceptor. A Cla Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI. +#### Making It Permanent at Login + +`lite up` holds the patch only for as long as it runs. To wire Claude Code up once and leave it that way, pass `--config-claude` to `lite login`: + +```bash +lite --base-url https://your-proxy.example.com login --config-claude +``` + +It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. + +Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. + +Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops. + ### QA Complexity-Based Auto-Routing Against Your Real Proxy `lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 1cac515f9f2..0a0bcf80ee5 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -16,6 +16,12 @@ from typing_extensions import NotRequired, TypedDict from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh +from .claude_settings import ( + CLAUDE_SETTINGS_PATH, + SETTINGS_FILE_OWNERS, + ClaudeSettingsError, + write_claude_settings, +) from .private_json import write_private_json @@ -629,9 +635,28 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: return None +def _configure_claude_code(base_url: str) -> None: + """Point Claude Code at base_url by patching ~/.claude/settings.json.""" + try: + write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + except ClaudeSettingsError as e: + raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") + click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.") + click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.") + + @click.command(name="login") +@click.option( + "--config-claude", + is_flag=True, + default=False, + help=( + "After logging in, update ~/.claude/settings.json so Claude Code routes through this proxy. " + "Unrelated settings are preserved." + ), +) @click.pass_context -def login(ctx: click.Context): +def login(ctx: click.Context, config_claude: bool): """Login to LiteLLM proxy using SSO authentication""" from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER from litellm.proxy.client.cli.interface import show_commands @@ -683,6 +708,9 @@ def login(ctx: click.Context): click.echo(f"JWT Token: {api_key[:20]}...") click.echo("You can now use the CLI without specifying --api-key") + if config_claude: + _configure_claude_code(base_url) + # Show available commands after successful login click.echo("\n" + "=" * 60) show_commands() @@ -698,6 +726,10 @@ def login(ctx: click.Context): except KeyboardInterrupt: click.echo("\nAuthentication cancelled by user.") return + except click.ClickException: + # Login itself already succeeded; only the post-login step failed, so this + # must not be relabelled as an authentication failure by the handler below. + raise except Exception as e: click.echo(f"Authentication failed: {e}") return diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 09e5a53b92f..26d45138a27 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -10,11 +10,16 @@ import click import yaml from pydantic import JsonValue, TypeAdapter, ValidationError -from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup +from ..claude_settings import ( + AUTOROUTE_BACKUP_PATH, + CLAUDE_SETTINGS_PATH, + ClaudeSettingsError, + load_json_or_empty, +) from ..up import BackupRecord as ClaudeBackupRecord +from ..up import restore_claude_settings, write_backup from .config import master_key_from_config from .process import ( - AUTOROUTE_DIR, CONFIG_PATH, DEFAULT_AUTOROUTE_PORT, LOG_PATH, @@ -35,8 +40,6 @@ from .process import ( from .settings import merge_claude_settings_static_token from .wizard import run_configure_wizard -AUTOROUTE_BACKUP_PATH: Final = AUTOROUTE_DIR / "claude_settings_backup.json" - _GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) @@ -108,7 +111,7 @@ def up(port: int) -> None: try: existing_pid: Final = read_pid_record() - except UpError as e: + except ClaudeSettingsError as e: raise click.ClickException(str(e)) if existing_pid is not None and is_running(existing_pid.pid): raise click.ClickException( @@ -157,7 +160,7 @@ def up(port: int) -> None: CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) with secure_create(CLAUDE_SETTINGS_PATH) as f: json.dump(merged, f, indent=2) - except UpError as e: + except ClaudeSettingsError as e: terminate(process.pid) clear_pid_record() raise click.ClickException(str(e)) @@ -175,7 +178,7 @@ def up(port: int) -> None: clear_pid_record() try: restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) - except UpError as e: + except ClaudeSettingsError as e: # Runs from atexit/a signal handler too, outside Click's own exception # handling -- raising here would only produce an unhandled-exception # warning on stderr, not a clean message. @@ -207,7 +210,7 @@ def down() -> None: """Restore Claude Code settings and stop a leftover ephemeral proxy, if any""" try: record: PidRecord | None = read_pid_record() - except UpError as e: + except ClaudeSettingsError as e: # down is the crash-recovery path -- a corrupt pid record must not block it; clear the # unusable record and keep going rather than leaving the user with no way to clean up. click.echo(f"{e} Clearing it and continuing cleanup.", err=True) @@ -219,7 +222,7 @@ def down() -> None: try: restored: Final = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) - except UpError as e: + except ClaudeSettingsError as e: raise click.ClickException(str(e)) if restored is None: click.echo("Nothing to restore.") diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py new file mode 100644 index 00000000000..e9a6a25a064 --- /dev/null +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -0,0 +1,155 @@ +"""Shared handling of Claude Code's ~/.claude/settings.json. + +`lite up` patches this file temporarily and restores it on exit; `lite login +--config-claude` patches it persistently. Both need the same merge and the same +apiKeyHelper command, and `up` already imports from `auth`, so the shared parts +live here rather than in either command module. +""" + +import shlex +import shutil +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from .private_json import write_private_json + +ENV_KEY: Final = "env" +API_KEY_HELPER_KEY: Final = "apiKeyHelper" +ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" + +CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" +BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" +AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" + + +@dataclass(frozen=True, slots=True) +class SettingsFileOwner: + """A command that takes temporary ownership of CLAUDE_SETTINGS_PATH and restores it later.""" + + backup_path: Path + start_command: str + stop_command: str + + +SETTINGS_FILE_OWNERS: Final = ( + SettingsFileOwner(BACKUP_PATH, "lite up", "lite down"), + SettingsFileOwner(AUTOROUTE_BACKUP_PATH, "lite autoroute up", "lite autoroute down"), +) + +_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) + + +class ClaudeSettingsError(Exception): + """Raised for any user-actionable failure while reading or writing Claude Code settings.""" + + +def load_json_or_empty(path: Path) -> dict[str, JsonValue]: + try: + content: Final = path.read_bytes() if path.exists() else b"" + except OSError as e: + raise ClaudeSettingsError(f"Could not read {path}: {e}") from e + if not content.strip(): + return {} + try: + return _SETTINGS_ADAPTER.validate_json(content) + except ValidationError: + raise ClaudeSettingsError( + f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely." + ) + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str +) -> dict[str, JsonValue]: + """Return a new settings dict wired to route Claude Code through the proxy. + + Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a + stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued + token (same reasoning as build_agent_env in agents.py). Every other key is + preserved untouched. + """ + raw_env: Final = settings.get(ENV_KEY, {}) + base_env: Final = raw_env if isinstance(raw_env, dict) else {} + env: Final = { + **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, + ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + } + return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} + + +def resolve_api_key_helper(base_url: str) -> str: + """Build the shell command Claude Code should run for its apiKeyHelper. + + Resolves `lite` to an absolute path so the helper works regardless of the + PATH visible to whatever subprocess Claude Code spawns it from. Passing + --base-url explicitly (rather than relying on the bare invocation Claude + Code would otherwise use) makes `print-token` enforce that the cached + token was actually issued for this proxy -- without it, a token minted + for a different, previously-logged-into proxy would be handed to + whichever server the settings currently point at. + + --base-url belongs to the top-level `lite` group, so it has to precede the + subcommand; click rejects it outright after `print-token`. + """ + lite_path: Final = shutil.which("lite") + if lite_path is None: + raise ClaudeSettingsError( + "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it." + ) + return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token" + + +def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: + """Persistently point Claude Code at base_url, preserving every unrelated setting. + + Refuses while any owner holds a backup: each restores its backup when it + stops, which would silently undo this write. + """ + for owner in owners: + if owner.backup_path.exists(): + raise ClaudeSettingsError( + f"`{owner.start_command}` is currently managing {settings_path} (backup at " + f"{owner.backup_path}) and will restore it when it stops. " + f"Run `{owner.stop_command}` first, then retry." + ) + normalized_base_url: Final = base_url.rstrip("/") + api_key_helper: Final = resolve_api_key_helper(normalized_base_url) + existing: Final = load_json_or_empty(settings_path) + raw_env: Final = existing.get(ENV_KEY) + if raw_env is not None and not isinstance(raw_env, dict): + raise ClaudeSettingsError( + f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. ' + "Fix or remove it, then retry." + ) + merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper) + # os.replace() swaps the symlink itself for a regular file, silently detaching a + # settings.json that is symlinked into a dotfiles repo. There is no backup to undo + # that here, unlike `lite up`, so write through to the link's target instead. + target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path + try: + write_private_json(str(target), merged) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {target}: {e}") from e + + +__all__ = ( + "ANTHROPIC_API_KEY_KEY", + "ANTHROPIC_BASE_URL_KEY", + "API_KEY_HELPER_KEY", + "AUTOROUTE_BACKUP_PATH", + "BACKUP_PATH", + "CLAUDE_SETTINGS_PATH", + "ENV_KEY", + "SETTINGS_FILE_OWNERS", + "ClaudeSettingsError", + "SettingsFileOwner", + "load_json_or_empty", + "merge_claude_settings", + "resolve_api_key_helper", + "write_claude_settings", +) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index 7023241cf06..dd266b4afa1 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -2,12 +2,10 @@ import atexit import contextlib import json import os -import shlex -import shutil import signal import sys import threading -from collections.abc import Iterator, Mapping +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from types import FrameType @@ -20,17 +18,17 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh from .agents import AgentRunError, resolve_api_key, verify_proxy_key from .auth import load_token, login - -ENV_KEY: Final = "env" -API_KEY_HELPER_KEY: Final = "apiKeyHelper" -ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" -ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" - -CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" -BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" +from .claude_settings import ( + BACKUP_PATH, + CLAUDE_SETTINGS_PATH, + ClaudeSettingsError, + load_json_or_empty, + merge_claude_settings, + resolve_api_key_helper, +) -class UpError(Exception): +class UpError(ClaudeSettingsError): """Raised for any user-actionable failure while starting/stopping interception.""" @@ -42,40 +40,9 @@ class BackupRecord: content: dict[str, JsonValue] | None -_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) _BACKUP_RECORD_ADAPTER: Final = TypeAdapter(BackupRecord) -def load_json_or_empty(path: Path) -> dict[str, JsonValue]: - if not path.exists(): - return {} - with open(path, "r") as f: - content: Final = f.read() - if not content.strip(): - return {} - try: - return _SETTINGS_ADAPTER.validate_json(content) - except ValidationError: - raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.") - - -def merge_claude_settings( - settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to route Claude Code through the proxy. - - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). Every other key is - preserved untouched. - """ - raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")} - env.pop(ANTHROPIC_API_KEY_KEY, None) - return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} - - @contextlib.contextmanager def secure_create(path: Path) -> Iterator[IO[str]]: """Open path for writing with mode 0600 fixed up before any content is written. @@ -136,26 +103,6 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path return record -def resolve_api_key_helper(base_url: str) -> str: - """Build the shell command Claude Code should run for its apiKeyHelper. - - Resolves `lite` to an absolute path so the helper works regardless of the - PATH visible to whatever subprocess Claude Code spawns it from. Passing - --base-url explicitly (rather than relying on the bare invocation Claude - Code would otherwise use) makes `print-token` enforce that the cached - token was actually issued for this proxy -- without it, a token minted - for a different, previously-logged-into proxy would be handed to - whichever server `up` currently points at. - """ - lite_path: Final = shutil.which("lite") - if lite_path is None: - raise UpError( - "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs " - "an absolute path to it, so `lite up` cannot continue." - ) - return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}" - - def _ensure_fresh_login(ctx: click.Context) -> None: base_url: Final = ctx.obj["base_url"].rstrip("/") token_data = load_token() @@ -224,7 +171,7 @@ def up(ctx: click.Context) -> None: merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper) with open(CLAUDE_SETTINGS_PATH, "w") as f: json.dump(merged, f, indent=2) - except (AgentRunError, UpError) as e: + except (AgentRunError, ClaudeSettingsError) as e: raise click.ClickException(str(e)) click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}") @@ -241,7 +188,7 @@ def up(ctx: click.Context) -> None: return try: _restore_and_report() - except UpError as e: + except ClaudeSettingsError as e: # Runs from atexit/a signal handler, outside Click's own exception # handling -- raising here would only produce an unhandled-exception # warning on stderr, not a clean message. @@ -264,7 +211,7 @@ def down() -> None: """ try: _restore_and_report() - except UpError as e: + except ClaudeSettingsError as e: raise click.ClickException(str(e)) @@ -272,6 +219,7 @@ __all__ = [ "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", "BackupRecord", + "ClaudeSettingsError", "UpError", "down", "load_json_or_empty", diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index f0aa49ff123..59048067674 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -25,6 +25,7 @@ from litellm.proxy.client.cli.commands.auth import ( save_token, whoami, ) +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner def _mock_cli_sso_start_response( @@ -267,7 +268,7 @@ class TestTokenUtilities: def test_load_token_io_error(self): """Test loading token with IO error""" with ( - patch("builtins.open", side_effect=IOError("Permission denied")), + patch("builtins.open", side_effect=OSError("Permission denied")), patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, patch("os.path.exists", return_value=True), ): @@ -1029,3 +1030,83 @@ class TestSaveTokenPrivateWrite: assert json.loads(token_file.read_text()) == {"key": "sk-original", "timestamp": 1234567890} assert list(token_file.parent.glob(".tmp-*")) == [] + + +class TestLoginConfigClaude: + """`lite login --config-claude` wiring into ~/.claude/settings.json""" + + def setup_method(self): + self.runner = CliRunner() + + def _run_login(self, tmp_path, args, base_url="https://test.example.com"): + settings_path = tmp_path / "claude" / "settings.json" + backup_path = tmp_path / "claude_settings_backup.json" + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + with ( + patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), + patch("requests.get", return_value=poll_response), + patch("litellm.proxy.client.cli.commands.auth.save_token"), + patch("litellm.proxy.client.cli.interface.show_commands"), + patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), + patch( + "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", + (SettingsFileOwner(backup_path, "lite up", "lite down"),), + ), + patch( + "litellm.proxy.client.cli.commands.claude_settings.shutil.which", + return_value="/usr/local/bin/lite", + ), + ): + result = self.runner.invoke(login, args, obj={"base_url": base_url}) + return result, settings_path, backup_path + + def test_default_login_does_not_touch_claude_settings(self, tmp_path): + result, settings_path, _backup_path = self._run_login(tmp_path, []) + + assert result.exit_code == 0 + assert "Login successful!" in result.output + assert not settings_path.exists() + assert "Configured Claude Code" not in result.output + + def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path): + result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + + assert result.exit_code == 0 + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" + assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" + assert "Configured Claude Code" in result.output + + def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path): + settings_path = tmp_path / "claude" / "settings.json" + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}})) + + result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + + assert result.exit_code == 0 + written = json.loads(settings_path.read_text()) + assert written["theme"] == "dark" + assert written["env"]["KEEP"] == "me" + + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path): + settings_path = tmp_path / "claude" / "settings.json" + settings_path.parent.mkdir(parents=True) + settings_path.write_text("not json at all {{{") + + result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + + assert result.exit_code != 0 + assert "Login successful!" in result.output + assert "could not configure Claude Code" in result.output + assert "invalid JSON" in result.output + assert "Authentication failed" not in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py new file mode 100644 index 00000000000..bc9744eb410 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -0,0 +1,269 @@ +import json +import shlex +import stat +import time +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands.claude_settings import ( + AUTOROUTE_BACKUP_PATH, + BACKUP_PATH, + SETTINGS_FILE_OWNERS, + ClaudeSettingsError, + SettingsFileOwner, + resolve_api_key_helper, + write_claude_settings, +) + + +def _owners(*backup_paths): + """Stand-in owners for the real `lite up` / `lite autoroute up` registry.""" + return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths) + +CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" +AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" + + +@pytest.fixture +def paths(tmp_path): + return tmp_path / "claude" / "settings.json", tmp_path / "backup.json" + + +@pytest.fixture +def lite_on_path(): + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): + yield + + +class TestWriteClaudeSettings: + def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path): + settings_path, backup_path = paths + assert not settings_path.parent.exists() + + write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path)) + + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + + def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text( + json.dumps( + { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"SOME_OTHER_VAR": "keep-me", "ANTHROPIC_BASE_URL": "https://old.example.com"}, + "apiKeyHelper": "old-helper", + } + ) + ) + + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + written = json.loads(settings_path.read_text()) + assert written["theme"] == "dark" + assert written["permissions"] == {"allow": ["Bash"]} + assert written["env"]["SOME_OTHER_VAR"] == "keep-me" + assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["apiKeyHelper"] != "old-helper" + + def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path): + settings_path, backup_path = paths + + write_claude_settings("https://first.example.com", settings_path, _owners(backup_path)) + write_claude_settings("https://second.example.com", settings_path, _owners(backup_path)) + + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com" + assert "second.example.com" in written["apiKeyHelper"] + assert "first.example.com" not in written["apiKeyHelper"] + + def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path): + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}})) + + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"] + + def test_written_file_is_owner_only(self, paths, lite_on_path): + settings_path, backup_path = paths + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600 + + def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path): + settings_path, backup_path = paths + backup_path.write_text("{}") + + with pytest.raises(ClaudeSettingsError, match="lite down"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + assert not settings_path.exists() + + def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths, lite_on_path): + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text("not json at all {{{") + + with pytest.raises(ClaudeSettingsError, match="invalid JSON"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + assert settings_path.read_text() == "not json at all {{{" + + def test_reports_an_actionable_error_when_lite_is_not_on_path(self, paths): + settings_path, backup_path = paths + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): + with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + assert not settings_path.exists() + + def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths, lite_on_path): + """Bytes that are not valid UTF-8 must not escape as UnicodeDecodeError. + + UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch + is easy to miss; login's broad `except Exception` would then relabel it as + an authentication failure and exit 0. + """ + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_bytes(b'{"theme": "\xff\xfe"}') + + with pytest.raises(ClaudeSettingsError, match="invalid JSON"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path): + """An unreadable settings file must not surface as "Authentication failed". + + login wraps the whole flow in a broad `except Exception`, so any OSError + escaping this function gets relabelled as an auth failure and sends the + user looking at their SSO config instead of at file permissions. + """ + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.mkdir() + + with pytest.raises(ClaudeSettingsError, match="Could not read"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path): + settings_path, backup_path = paths + with patch( + f"{CLAUDE_SETTINGS_MODULE}.write_private_json", + side_effect=OSError("Read-only file system"), + ): + with pytest.raises(ClaudeSettingsError, match="Read-only file system"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + +class TestApiKeyHelperIsActuallyInvocable: + """The helper string is executed verbatim by Claude Code, so it has to parse. + + Asserting only on its text is what let a malformed command (`--base-url`, a + top-level group option, placed after the `print-token` subcommand) ship: click + rejects it with "No such option" and every Claude Code request loses its token. + """ + + def _helper_args(self, base_url): + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): + return shlex.split(resolve_api_key_helper(base_url))[1:] + + def test_the_generated_command_parses(self): + result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) + + assert "No such option" not in result.output + assert result.exit_code != 2 + + def test_the_generated_command_reaches_print_token(self): + with patch(f"{AUTH_MODULE}.load_token", return_value=None): + result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) + + assert "Not authenticated" in result.output + + def test_the_generated_command_carries_the_base_url_through(self): + stale = { + "base_url": "http://other-proxy.example.com", + "key": "sk-stale", + "timestamp": time.time(), + } + with patch(f"{AUTH_MODULE}.load_token", return_value=stale): + result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) + + assert "Not authenticated for this server" in result.output + + +class TestConflictingOwnersOfTheSettingsFile: + """Both `lite up` and `lite autoroute up` restore a backup when they stop. + + Guarding only one of them leaves the other free to silently revert this + write, which is the exact hazard the guard exists to prevent. + """ + + def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path, lite_on_path): + settings_path = tmp_path / "claude" / "settings.json" + + for index, owner in enumerate(SETTINGS_FILE_OWNERS): + backup = tmp_path / f"backup-{index}.json" + backup.write_text("{}") + stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command) + with pytest.raises(ClaudeSettingsError, match="currently managing"): + write_claude_settings("https://proxy.example.com", settings_path, (stand_in,)) + backup.unlink() + assert not settings_path.exists() + + def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path, lite_on_path): + settings_path = tmp_path / "claude" / "settings.json" + backup = tmp_path / "auto.json" + backup.write_text("{}") + autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down") + + with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"): + write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"): + write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + + def test_the_registry_matches_the_paths_the_commands_actually_use(self): + """A second definition of the autoroute dir must not drift from this one.""" + from litellm.proxy.client.cli.commands.autoroute.process import AUTOROUTE_DIR + + assert AUTOROUTE_BACKUP_PATH == AUTOROUTE_DIR / "claude_settings_backup.json" + assert {o.backup_path for o in SETTINGS_FILE_OWNERS} == {BACKUP_PATH, AUTOROUTE_BACKUP_PATH} + assert {o.stop_command for o in SETTINGS_FILE_OWNERS} == {"lite down", "lite autoroute down"} + + +class TestDoesNotDestroyUserOwnedStructure: + def test_writes_through_a_symlinked_settings_file(self, tmp_path, lite_on_path): + """os.replace() swaps the symlink for a regular file, detaching a dotfiles repo. + + There is no backup here to undo that, so the link must survive and its + target must be the thing that gets updated. + """ + real = tmp_path / "dotfiles" / "settings.json" + real.parent.mkdir() + real.write_text(json.dumps({"theme": "dark"})) + link = tmp_path / "claude" / "settings.json" + link.parent.mkdir() + link.symlink_to(real) + + write_claude_settings("https://proxy.example.com", link, ()) + + assert link.is_symlink() + assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert json.loads(real.read_text())["theme"] == "dark" + + def test_refuses_rather_than_discarding_a_non_object_env(self, paths, lite_on_path): + """merge coerces a non-dict env to {}; that is silent data loss on a persistent write.""" + settings_path, backup_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"})) + + with pytest.raises(ClaudeSettingsError, match="non-object"): + write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + + assert json.loads(settings_path.read_text())["env"] == "not-an-object" diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 1b182553644..51de0dcf11d 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -10,6 +10,7 @@ from click.testing import CliRunner from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, @@ -25,6 +26,7 @@ from litellm.proxy.client.cli.commands.up import ( ) UP_MODULE = "litellm.proxy.client.cli.commands.up" +AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" def _patch_paths(monkeypatch, tmp_path): @@ -92,13 +94,13 @@ class TestLoadJsonOrEmpty: def test_raises_clean_error_on_invalid_json(self, tmp_path): path = tmp_path / "settings.json" path.write_text("not json at all {{{") - with pytest.raises(UpError, match="invalid JSON"): + with pytest.raises(ClaudeSettingsError, match="invalid JSON"): load_json_or_empty(path) def test_raises_clean_error_on_non_object_root(self, tmp_path): path = tmp_path / "settings.json" path.write_text(json.dumps([1, 2, 3])) - with pytest.raises(UpError, match="invalid JSON"): + with pytest.raises(ClaudeSettingsError, match="invalid JSON"): load_json_or_empty(path) @@ -201,16 +203,16 @@ class TestResolveApiKeyHelper: def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch): monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") helper = resolve_api_key_helper("http://localhost:4000") - assert helper == "/usr/local/bin/lite auth print-token --base-url http://localhost:4000" + assert helper == "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token" def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch): monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") helper = resolve_api_key_helper("http://example.com/path; rm -rf /") - assert helper == "/usr/local/bin/lite auth print-token --base-url 'http://example.com/path; rm -rf /'" + assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" def test_raises_when_lite_not_on_path(self, monkeypatch): monkeypatch.setattr(shutil, "which", lambda name: None) - with pytest.raises(UpError, match="Could not find `lite`"): + with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): resolve_api_key_helper("http://localhost:4000") @@ -396,3 +398,50 @@ class TestDownCommand: assert result.exit_code != 0 assert result.exception is None or isinstance(result.exception, SystemExit) assert "invalid or unexpected JSON" in result.output + + +class TestUpCanInvokeTheRealLoginCommand: + """`lite up` calls ctx.invoke(login) on the real command object. + + Every other test in this file monkeypatches `up_module.login` with a fake, so + none of them would notice a login parameter that ctx.invoke cannot supply. + """ + + def test_ctx_invoke_supplies_every_login_parameter(self): + from litellm.proxy.client.cli.commands.auth import login as real_login + + reached = [] + + @click.command() + @click.pass_context + def driver(ctx): + ctx.obj = {"base_url": "http://127.0.0.1:9"} + ctx.invoke(real_login) + + with patch( + f"{AUTH_MODULE}._start_cli_sso_flow", + side_effect=lambda base_url: reached.append(base_url) or RuntimeError("stop"), + ): + result = CliRunner().invoke(driver, [], standalone_mode=False) + + assert not isinstance(result.exception, TypeError), result.exception + assert reached == ["http://127.0.0.1:9"] + + def test_ctx_invoke_leaves_claude_settings_alone(self, tmp_path): + from litellm.proxy.client.cli.commands.auth import login as real_login + + settings_path = tmp_path / "settings.json" + + @click.command() + @click.pass_context + def driver(ctx): + ctx.obj = {"base_url": "http://127.0.0.1:9"} + ctx.invoke(real_login) + + with ( + patch(f"{AUTH_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), + patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")), + ): + CliRunner().invoke(driver, [], standalone_mode=False) + + assert not settings_path.exists() From 74b279bc44e12c4b73ac9906c834dd677724a967 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 15:33:29 -0700 Subject: [PATCH 58/88] fix(auth): resolve bare model names against wildcard deployments in model access groups (#37492) * fix(auth): resolve bare model names against wildcard deployments in model access groups * test(e2e): cover model access group permission checks on keys and teams --- litellm/router.py | 6 +- .../access_control/access_control_client.py | 82 +++++- .../test_model_access_group_e2e.py | 277 ++++++++++++++++++ tests/e2e/coverage_registry/other.yaml | 5 + tests/e2e/models.py | 3 + tests/e2e/proxy_client.py | 29 +- .../proxy/auth/test_auth_checks.py | 165 +++++++++++ 7 files changed, 553 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/access_control/test_model_access_group_e2e.py diff --git a/litellm/router.py b/litellm/router.py index 00c6c4c8b6f..b25b4f92467 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10253,11 +10253,13 @@ class Router: returned_models.extend(self.get_model_list_from_routing_groups(model_name=model_name)) if len(returned_models) == 0: # check if wildcard route - potential_wildcard_models: Final = self.pattern_router.route(model_name) or [] + potential_wildcard_models: Final = self.pattern_router.get_deployments_by_pattern(model=model_name or "") ## check for team-specific wildcard models if team_id is not None and team_id in self.team_pattern_routers: - potential_team_only_wildcard_models: Final = self.team_pattern_routers[team_id].route(model_name) or [] + potential_team_only_wildcard_models: Final = self.team_pattern_routers[ + team_id + ].get_deployments_by_pattern(model=model_name or "") potential_wildcard_models.extend(potential_team_only_wildcard_models) if model_name is not None and potential_wildcard_models is not None: diff --git a/tests/e2e/access_control/access_control_client.py b/tests/e2e/access_control/access_control_client.py index 7ace036f433..634a96bb0bd 100644 --- a/tests/e2e/access_control/access_control_client.py +++ b/tests/e2e/access_control/access_control_client.py @@ -2,12 +2,13 @@ from __future__ import annotations +import time from dataclasses import dataclass from pydantic import BaseModel, ValidationError from proxy_client import ProxyClient -from e2e_http import StreamingResponse +from e2e_http import NoBody, StreamingResponse, is_ok, unwrap from models import ( ChatBody, ChatMessage, @@ -15,9 +16,16 @@ from models import ( LiteLLMParamsBody, ModelInfoBody, ModelNewBody, + TeamDeleteBody, + TeamInfoParams, + TeamInfoResponse, + TeamNewBody, + TeamNewResponse, + TeamUpdateBody, ) MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" +TEAM_MODEL_ACCESS_DENIED_MARKER = "team_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @@ -31,6 +39,14 @@ class ApiErrorEnvelope(BaseModel): error: ApiErrorDetail +class AccessGroupInfoResponse(BaseModel): + """GET /access_group/{name}/info: the deployments a model access group grants.""" + + access_group: str + model_names: list[str] + deployment_count: int + + def error_envelope(body: str) -> ApiErrorEnvelope | None: """The OpenAI-shaped `{"error": {...}}` a client parses, or None if absent.""" try: @@ -51,15 +67,75 @@ class AccessControlClient: def delete_key(self, key: str) -> None: self.proxy.delete_key(key) - def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: + def chat_status( + self, key: str, model: str, content: str, max_completion_tokens: int | None = None + ) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", headers=self.proxy.transport.bearer(key), json=ChatBody( - model=model, messages=[ChatMessage(role="user", content=content)] + model=model, + messages=[ChatMessage(role="user", content=content)], + max_completion_tokens=max_completion_tokens, ), ) + def create_team(self, team_alias: str, models: list[str]) -> str: + team_id = unwrap( + self.proxy.transport.post( + "/team/new", + headers=self.proxy.transport.master, + json=TeamNewBody(team_alias=team_alias, models=models), + response_type=TeamNewResponse, + ) + ).team_id + self._await_team(team_id) + return team_id + + def set_team_models(self, team_id: str, team_alias: str, models: list[str]) -> None: + """Replace the team's allow-list. /model/new appends a team-scoped deployment's + public name to it, so a test that means to grant only an access group has to + put the allow-list back afterwards.""" + _ = unwrap( + self.proxy.transport.post( + "/team/update", + headers=self.proxy.transport.master, + json=TeamUpdateBody(team_id=team_id, team_alias=team_alias, models=models), + response_type=NoBody, + ) + ) + + def delete_team(self, team_id: str) -> None: + _ = self.proxy.transport.post( + "/team/delete", + headers=self.proxy.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def access_group_info(self, access_group: str) -> AccessGroupInfoResponse | None: + result = self.proxy.transport.get( + f"/access_group/{access_group}/info", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AccessGroupInfoResponse, + ) + return unwrap(result) if is_ok(result) else None + + def _await_team(self, team_id: str) -> None: + deadline = time.monotonic() + self.proxy.poll_timeout + while time.monotonic() < deadline: + result = self.proxy.transport.get( + "/team/info", + headers=self.proxy.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + if is_ok(result): + return + time.sleep(self.proxy.poll_interval) + raise AssertionError(f"/team/info never resolved team {team_id!r} created by /team/new") + def create_model_status(self, key: str, model_name: str) -> StreamingResponse: return self.proxy.transport.send( "/model/new", diff --git a/tests/e2e/access_control/test_model_access_group_e2e.py b/tests/e2e/access_control/test_model_access_group_e2e.py new file mode 100644 index 00000000000..5cc062ea096 --- /dev/null +++ b/tests/e2e/access_control/test_model_access_group_e2e.py @@ -0,0 +1,277 @@ +"""Live e2e: a model access group as the grant on a key and on a team. + +Whoever holds the group can call every deployment in it and nothing else, whether +the request names a deployment exactly, names a model that a wildcard deployment +in the group covers, or spells that model with its provider prefix. The bare-name +spelling is the LIT-5813 regression: the group-membership lookup skipped the +provider-prefix retry every other model-resolution path performs, so a group +holding `openai/gpt-5.4*` denied `gpt-5.4-nano` while allowing `openai/gpt-5.4-nano`. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from typing import Final + +import pytest + +from access_control_client import ( + AccessControlClient, + MODEL_ACCESS_DENIED_MARKER, + TEAM_MODEL_ACCESS_DENIED_MARKER, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ( + ChatResponse, + KeyGenerateBody, + LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, +) + +pytestmark = pytest.mark.e2e + +WILDCARD_PATTERN: Final = "openai/gpt-5.4*" +WILDCARD_BARE_MODEL: Final = "gpt-5.4-nano" +WILDCARD_PREFIXED_MODEL: Final = "openai/gpt-5.4-nano" +GROUP_BACKEND: Final = "openai/gpt-5.4-nano" +UNCOVERED_OPENAI_MODEL: Final = "gpt-5.2" + +TEAM_WILDCARD_PATTERN: Final = "openai/gpt-5.6*" +TEAM_WILDCARD_BARE_MODEL: Final = "gpt-5.6-luna" + +MAX_COMPLETION_TOKENS: Final = 16 +PROMPT: Final = "Reply with exactly: OK" + + +@dataclass(frozen=True, slots=True) +class GroupedDeployments: + """A wildcard deployment and an exactly-named one inside `access_group`, plus a + deployment left out of it.""" + + access_group: str + member_model: str + outsider_model: str + + +@dataclass(frozen=True, slots=True) +class TeamGrant: + """A team whose whole allow-list is `access_group`, holding one team-scoped + wildcard deployment, and a key that belongs to it.""" + + access_group: str + team_id: str + key: str + + +ModelSelector = Callable[[GroupedDeployments], str] + +ALLOWED: Final[tuple[tuple[str, ModelSelector], ...]] = ( + ("bare name the group's wildcard covers", lambda grouped: WILDCARD_BARE_MODEL), + ("provider-prefixed name the group's wildcard covers", lambda grouped: WILDCARD_PREFIXED_MODEL), + ("exactly-named deployment in the group", lambda grouped: grouped.member_model), +) + +DENIED: Final[tuple[tuple[str, ModelSelector], ...]] = ( + ("deployment outside the group", lambda grouped: grouped.outsider_model), + ("provider model outside the group's wildcard", lambda grouped: UNCOVERED_OPENAI_MODEL), + ("name no provider claims", lambda grouped: f"e2e-ag-unknown-{unique_marker()}"), +) + + +def _provider_key(env_var: str) -> str: + return os.environ.get(env_var) or f"os.environ/{env_var}" + + +def _grouped_model(model_name: str, backend: str, access_groups: list[str] | None) -> ModelNewBody: + return ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody(model=backend, api_key=_provider_key("OPENAI_API_KEY")), + model_info=ModelInfoBody(access_groups=access_groups), + ) + + +def _await_group_members(client: AccessControlClient, access_group: str, expected: frozenset[str]) -> None: + """The grant under test is the group's membership, so prove the proxy recorded it + before asserting on what the group lets through.""" + deadline = time.monotonic() + client.proxy.poll_timeout + listed: list[str] = [] + while time.monotonic() < deadline: + info = client.access_group_info(access_group) + listed = info.model_names if info is not None else [] + if expected.issubset(listed): + return + time.sleep(client.proxy.poll_interval) + pytest.fail( + f"/access_group/{access_group}/info never listed {sorted(expected)} as members; last read {listed}" + ) + + +def _await_team_allowlist(client: AccessControlClient, grant_key: str, access_group: str) -> None: + """Registering a team-scoped deployment appends its public name to the team's + allow-list, and a wildcard sitting there directly would grant the model under test + on its own. Poll a denial until the message enumerates the allow-list the test + means to exercise: the group, and nothing else.""" + allowlist: Final = f"models=['{access_group}']" + deadline = time.monotonic() + client.proxy.poll_timeout + body = "" + while time.monotonic() < deadline: + body = client.chat_status( + grant_key, UNCOVERED_OPENAI_MODEL, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS + ).body + if allowlist in body: + return + time.sleep(client.proxy.poll_interval) + pytest.fail(f"the team's allow-list never settled to {allowlist}; last denial read {body[:300]}") + + +@pytest.fixture(scope="module") +def grouped(client: AccessControlClient) -> Iterator[GroupedDeployments]: + marker: Final = unique_marker() + deployments: Final = GroupedDeployments( + access_group=f"e2e-ag-{marker}", + member_model=f"e2e-ag-member-{marker}", + outsider_model=f"e2e-ag-outsider-{marker}", + ) + registrations: Final = ( + _grouped_model(WILDCARD_PATTERN, WILDCARD_PATTERN, [deployments.access_group]), + _grouped_model(deployments.member_model, GROUP_BACKEND, [deployments.access_group]), + _grouped_model(deployments.outsider_model, GROUP_BACKEND, None), + ) + created: Final = tuple(client.proxy.register_model(body) for body in registrations) + try: + _await_group_members( + client, + deployments.access_group, + frozenset({WILDCARD_PATTERN, deployments.member_model}), + ) + yield deployments + finally: + for model_id in created: + client.proxy.delete_model(model_id) + + +@pytest.fixture(scope="module") +def team_grant(client: AccessControlClient) -> Iterator[TeamGrant]: + marker: Final = unique_marker() + access_group: Final = f"e2e-agt-{marker}" + team_alias: Final = f"e2e-ag-team-{marker}" + team_id: Final = client.create_team(team_alias, [access_group]) + key: Final = client.proxy.generate_key(KeyGenerateBody(models=[], team_id=team_id)) + model_id: Final = client.proxy.register_model( + ModelNewBody( + model_name=TEAM_WILDCARD_PATTERN, + litellm_params=LiteLLMParamsBody( + model=TEAM_WILDCARD_PATTERN, api_key=_provider_key("OPENAI_API_KEY") + ), + model_info=ModelInfoBody(team_id=team_id, access_groups=[access_group]), + ), + listed_for=key, + ) + client.set_team_models(team_id, team_alias, [access_group]) + try: + _await_team_allowlist(client, key, access_group) + yield TeamGrant(access_group=access_group, team_id=team_id, key=key) + finally: + client.proxy.delete_model(model_id) + client.proxy.delete_key(key) + client.delete_team(team_id) + + +class TestKeyScopedToAccessGroup: + @pytest.mark.covers( + "other.auth.model_access_group.wildcard_bare_name_allowed", + "other.auth.model_access_group.member_allowed", + ) + @pytest.mark.parametrize(("case", "select_model"), ALLOWED) + def test_group_grants_every_deployment_in_it( + self, + case: str, + select_model: ModelSelector, + client: AccessControlClient, + resources: ResourceManager, + grouped: GroupedDeployments, + ) -> None: + key = resources.key(models=[grouped.access_group]) + model = select_model(grouped) + + result = client.chat_status( + key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS + ) + + assert result.status_code == 200, ( + f"a key holding access group {grouped.access_group!r} must be able to call " + f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}" + ) + assert ChatResponse.model_validate_json(result.body).choices, ( + f"200 must carry a real completion, not an error envelope: {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.model_access_group.non_member_denied") + @pytest.mark.parametrize(("case", "select_model"), DENIED) + def test_group_grants_nothing_outside_it( + self, + case: str, + select_model: ModelSelector, + client: AccessControlClient, + resources: ResourceManager, + grouped: GroupedDeployments, + ) -> None: + key = resources.key(models=[grouped.access_group]) + model = select_model(grouped) + + result = client.chat_status( + key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS + ) + + assert result.status_code == 403, ( + f"a key holding only access group {grouped.access_group!r} must be denied 403 on " + f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}" + ) + assert MODEL_ACCESS_DENIED_MARKER in result.body, ( + f"403 body must be a key model-access denial, got: {result.body[:300]}" + ) + + +class TestTeamScopedToAccessGroup: + @pytest.mark.covers("other.auth.model_access_group.team_wildcard_bare_name_allowed") + def test_group_grants_the_teams_own_wildcard( + self, client: AccessControlClient, team_grant: TeamGrant + ) -> None: + result = client.chat_status( + team_grant.key, + TEAM_WILDCARD_BARE_MODEL, + f"{PROMPT} {unique_marker()}", + MAX_COMPLETION_TOKENS, + ) + + assert result.status_code == 200, ( + f"a team whose allow-list is access group {team_grant.access_group!r} must be able to " + f"call {TEAM_WILDCARD_BARE_MODEL!r} through its team-scoped {TEAM_WILDCARD_PATTERN!r} " + f"deployment, got {result.status_code}: {result.body[:300]}" + ) + assert ChatResponse.model_validate_json(result.body).choices, ( + f"200 must carry a real completion, not an error envelope: {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.model_access_group.team_non_member_denied") + def test_group_grants_the_team_nothing_outside_it( + self, client: AccessControlClient, team_grant: TeamGrant + ) -> None: + model = f"e2e-ag-unknown-{unique_marker()}" + + result = client.chat_status( + team_grant.key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS + ) + + assert result.status_code == 403, ( + f"a team holding only access group {team_grant.access_group!r} must be denied 403 on " + f"{model!r}, got {result.status_code}: {result.body[:300]}" + ) + assert TEAM_MODEL_ACCESS_DENIED_MARKER in result.body, ( + f"403 body must be a team model-access denial, got: {result.body[:300]}" + ) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index c7140a4503b..814ebae2e0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -12,6 +12,11 @@ - {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"} - {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"} - {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"} +- {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} +- {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} +- {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} +- {id: other.auth.model_access_group.team_wildcard_bare_name_allowed, module: other, tier: P1, area: auth, assertions: [team_wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "The same bare-name grant holds when the wildcard deployment is team-scoped and the team's allow-list is the group"} +- {id: other.auth.model_access_group.team_non_member_denied, module: other, tier: P1, area: auth, assertions: [team_non_member_denied], source: "auth_checks.py:3232", rationale: "A team-level group grant reaches nothing outside the group"} - {id: other.auth.virtual_key.route_permission_enforced, module: other, tier: P0, area: auth, assertions: [route_permission_enforced], source: "route_checks.py:89-151", rationale: "allowed_routes whitelist denies disallowed routes"} - {id: other.auth.virtual_key.route_group_allowed, module: other, tier: P1, area: auth, assertions: [route_group_allowed], source: "route_checks.py:106-128", rationale: "allowed_routes=[llm_api_routes] grants all LLM endpoints"} - {id: other.auth.passthrough.model_allowlist_enforced, module: other, tier: P1, area: auth, assertions: [model_allowlist_enforced], source: "route_checks.py:135-151", rationale: "Passthrough enforces per-key model allow-lists"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index df5cb841fad..619f4dcacfe 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -766,6 +766,8 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None + access_groups: list[str] | None = None + team_id: str | None = None class ModelNewBody(BaseModel): @@ -861,6 +863,7 @@ class TeamNewResponse(BaseModel): class TeamUpdateBody(BaseModel): team_id: str team_alias: str + models: list[str] | None = None class TeamInfoParams(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 843799ede6c..3cae337a5ff 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -278,7 +278,21 @@ class ProxyClient: mode: ModelMode | None = None, ) -> str: """Register a deployment under `model_name` and return its proxy-assigned - model_id, once the model is actually servable on the data plane. + model_id, once the model is actually servable on the data plane.""" + return self.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=litellm_params, + model_info=ModelInfoBody(mode=mode), + ) + ) + + def register_model(self, body: ModelNewBody, listed_for: str | None = None) -> str: + """`create_model` for deployments that carry more than a mode: access groups, + team scoping, a pinned id. `listed_for` is the virtual key whose /v1/models + view must list the deployment before it counts as servable, because a + team-scoped deployment is listed to its own team and to nobody else, master + key included; leave it unset for a proxy-wide model. /model/new is a control-plane route; the data plane (which serves /chat, /ocr, ...) only picks the new model up on its next DB reload, so a call @@ -296,25 +310,22 @@ class ProxyClient: self.transport.post( "/model/new", headers=self.transport.master, - json=ModelNewBody( - model_name=model_name, - litellm_params=litellm_params, - model_info=ModelInfoBody(mode=mode), - ), + json=body, response_type=ModelNewResponse, ) ).model_id written_at = time.monotonic() - self._await_model_servable(model_name) + self._await_model_servable(body.model_name, listed_for) settle_propagation(written_at) return model_id - def _await_model_servable(self, model_name: str) -> None: + def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until the data plane lists `model_name`, or fail at model_servable_timeout.""" + headers = self.transport.master if listed_for is None else self.transport.bearer(listed_for) outcome = await_servable( lambda poll_timeout: self.transport.get( "/v1/models", - headers=self.transport.master, + headers=headers, params=NoBody(), response_type=ModelsListResponse, timeout=poll_timeout, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 270f3eca0f9..1d1bd9ebf8a 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6392,3 +6392,168 @@ def test_is_user_proxy_admin_rejects_view_only_admin(): assert _is_user_proxy_admin(user_obj=viewer) is False assert _is_user_proxy_admin(user_obj=admin) is True assert _is_user_proxy_admin(user_obj=None) is False + + +def _make_wildcard_access_group_router(): + """ + `openai/*` tagged into an access group, plus an untagged `azure/*`, mirroring a + proxy that fronts a whole provider behind one wildcard deployment. + """ + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "fake"}, + "model_info": { + "id": "wildcard-openai", + "access_groups": ["default-models"], + }, + }, + { + "model_name": "azure/*", + "litellm_params": {"model": "azure/*", "api_key": "fake"}, + "model_info": {"id": "wildcard-azure"}, + }, + ] + ) + + +def test_can_object_call_model_access_group_wildcard_accepts_bare_model_name(): + """ + Regression: a key holding only the access group name was denied for `gpt-4o` + while `openai/gpt-4o` was allowed, because group membership resolved through the + pattern router's raw regex and skipped the `{provider}/{model}` retry that both + routing and the direct-wildcard grant already perform. + """ + from litellm.proxy.auth.auth_checks import _can_object_call_model + + router = _make_wildcard_access_group_router() + + assert ( + _can_object_call_model( + model="gpt-4o", + llm_router=router, + models=["default-models"], + object_type="key", + ) + is True + ) + + +def test_can_object_call_model_access_group_wildcard_accepts_prefixed_model_name(): + from litellm.proxy.auth.auth_checks import _can_object_call_model + + router = _make_wildcard_access_group_router() + + assert ( + _can_object_call_model( + model="openai/gpt-4o", + llm_router=router, + models=["default-models"], + object_type="key", + ) + is True + ) + + +@pytest.mark.parametrize( + "model", + [ + "totally-made-up-model-zzz", # no provider can be inferred + "azure/some-deployment", # wildcard exists but carries no access group + ], +) +def test_can_object_call_model_access_group_wildcard_does_not_over_grant(model): + """The bare-name retry must not turn an access group into a blanket grant.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import _can_object_call_model + + router = _make_wildcard_access_group_router() + + with pytest.raises(ProxyException): + _can_object_call_model( + model=model, + llm_router=router, + models=["default-models"], + object_type="key", + ) + + +def test_can_object_call_model_access_group_rejects_unconsumed_namespace(): + """ + `bedrockz/...` infers provider `bedrock` from a fragment of the name, so + re-prefixing would smuggle an unrecognized namespace through a `bedrock/*` group. + """ + from litellm import Router + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import _can_object_call_model + + router = Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": { + "id": "wildcard-bedrock", + "access_groups": ["bedrock-models"], + }, + } + ] + ) + + assert ( + _can_object_call_model( + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_router=router, + models=["bedrock-models"], + object_type="key", + ) + is True + ) + + with pytest.raises(ProxyException): + _can_object_call_model( + model="bedrockz/anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_router=router, + models=["bedrock-models"], + object_type="key", + ) + + +def test_can_object_call_model_team_scoped_wildcard_accepts_bare_model_name(): + """ + Same regression as the proxy-wide wildcard, but for a team-scoped deployment + whose public name is a wildcard: those live in a separate per-team pattern + index that needed the same `{provider}/{model}` retry. + """ + from litellm import Router + from litellm.proxy.auth.auth_checks import _can_object_call_model + + router = Router( + model_list=[ + { + "model_name": "openai/*_team-a_abc", + "litellm_params": {"model": "openai/*", "api_key": "fake"}, + "model_info": { + "id": "team-byok-wildcard", + "team_id": "team-a", + "team_public_model_name": "openai/*", + "access_groups": ["team-models"], + }, + } + ] + ) + + for model in ("gpt-4o", "openai/gpt-4o"): + assert ( + _can_object_call_model( + model=model, + llm_router=router, + models=["team-models"], + object_type="team", + team_id="team-a", + ) + is True + ) From b477d0967a4bf1f4a869c70ddf5a740f1762ae86 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:34:38 -0700 Subject: [PATCH 59/88] fix(proxy): lift standard_logging_object onto request_data before the logging object is popped --- litellm/proxy/utils.py | 52 ++++++++++--------- tests/test_litellm/proxy/test_proxy_utils.py | 53 ++++++++++++++++++++ 2 files changed, 82 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a743526e975..58ac7882cb8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -517,6 +517,34 @@ def _failure_usage_to_lift( return estimated_usage, 0.0 +_EMPTY_LIFT: Final = MappingProxyType({}) + + +def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: + """Failure-path callbacks run after ``litellm_logging_obj`` is popped from + request_data (it is not serialisable), so the caller merges these fields + onto request_data first: the first-handoff instant for preprocessing + latency, recovered or estimated usage for token counts, and the standard + logging object for deployment attribution on failed-request spend logs.""" + _logging_obj: Final = request_data.get("litellm_logging_obj") + if _logging_obj is None: + return _EMPTY_LIFT + _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) + _first_handoff: Final = _model_call_details.get("first_api_call_start_time") + _usage_to_lift: Final = _failure_usage_to_lift( + model_call_details=_model_call_details, + request_body=request_data, + dispatched=_first_handoff is not None, + ) + _entries: Final = ( + ("first_api_call_start_time", _first_handoff), + ("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]), + ("response_cost", None if _usage_to_lift is None else _usage_to_lift[1]), + ("standard_logging_object", _model_call_details.get("standard_logging_object")), + ) + return MappingProxyType({key: value for key, value in _entries if value is not None}) + + @dataclass(frozen=True) class _CallbackCapabilities: """Cached per-hook capability flags derived from ``litellm.callbacks``. @@ -2294,29 +2322,7 @@ class ProxyLogging: original_exception=original_exception, ) - # Lift the first-handoff instant onto request_data (top-level - # internal key, not metadata) so failure-path callbacks can still - # compute preprocessing latency after the logging object is popped. - _logging_obj: Final = request_data.get("litellm_logging_obj") - if _logging_obj is not None: - _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) - _first_handoff: Final = _model_call_details.get("first_api_call_start_time") - if _first_handoff is not None: - request_data["first_api_call_start_time"] = _first_handoff - - # Lift recovered partial-stream usage, or an estimated input-side - # usage for a dispatched failure, onto request_data so the - # failure-path spend callbacks (which run after the logging object - # is popped) record real token counts instead of zero. - _usage_to_lift: Final = _failure_usage_to_lift( - model_call_details=_model_call_details, - request_body=request_data, - dispatched=_first_handoff is not None, - ) - if _usage_to_lift is not None: - _lifted_usage, _lifted_cost = _usage_to_lift - request_data["combined_usage_object"] = _lifted_usage - request_data["response_cost"] = _lifted_cost + request_data.update(_failure_fields_to_lift(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index d6cf0e30139..c80130b44da 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -478,6 +478,59 @@ class TestPostCallFailureHookLiftsRecoveredPartialSpend: assert "response_cost" not in request_data +class TestPostCallFailureHookLiftsStandardLoggingObject: + """Failure callbacks read standard_logging_object from request_data, but + post_call_failure_hook pops litellm_logging_obj before they run. The hook + must lift the logging obj's standard_logging_object onto request_data so + failed-request spend logs keep deployment attribution (LIT-5795). + """ + + async def _run(self, request_data): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + @pytest.mark.asyncio + async def test_lifts_standard_logging_object(self): + sl_object = {"model_id": "mid-123", "model_group": "group-x"} + logging_obj = MagicMock() + logging_obj.model_call_details = {"standard_logging_object": sl_object} + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + assert request_data["standard_logging_object"] is sl_object + assert "litellm_logging_obj" not in request_data + + @pytest.mark.asyncio + async def test_logging_obj_value_overwrites_preexisting_key(self): + authoritative = {"model_id": "from-logging-obj"} + logging_obj = MagicMock() + logging_obj.model_call_details = {"standard_logging_object": authoritative} + request_data = { + "litellm_logging_obj": logging_obj, + "standard_logging_object": {"model_id": "client-injected"}, + "metadata": {}, + } + await self._run(request_data) + assert request_data["standard_logging_object"] is authoritative + + @pytest.mark.asyncio + async def test_no_standard_logging_object_is_noop(self): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + assert "standard_logging_object" not in request_data + + class TestPostCallFailureHookEstimatesDispatchedInputTokens: """A non-stream request that failed after dispatch (timeout, provider error) consumed provider-billed input tokens but recovered no usage. From 9bb54839919c27477d7cc88ce4d1caef3b08ce59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:37:44 -0700 Subject: [PATCH 60/88] fix(proxy): apply db-backed max_batch_file_size_mb on config reload --- litellm/proxy/proxy_server.py | 3 ++ .../proxy/proxy_server/test_proxy_config.py | 37 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f08c9887a92..26bbf767ed2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6316,6 +6316,9 @@ class ProxyConfig: if "global_max_parallel_requests" in _general_settings: general_settings["global_max_parallel_requests"] = _general_settings["global_max_parallel_requests"] + if "max_batch_file_size_mb" not in self._yaml_general_settings_keys: + general_settings["max_batch_file_size_mb"] = _general_settings.get("max_batch_file_size_mb") + ## ALERTING ARGS ## if "alerting_args" in _general_settings: general_settings["alerting_args"] = _general_settings["alerting_args"] diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index f31f67c317a..37a03617a65 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2442,6 +2442,43 @@ async def test_ProxyConfig__update_general_settings_updates_max_parallel(monkeyp } +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_applies_db_max_batch_file_size_mb(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + pc = ProxyConfig() + await pc._update_general_settings({"max_batch_file_size_mb": 5}) + from litellm.proxy import proxy_server as ps + + assert ps.general_settings.get("max_batch_file_size_mb") == 5 + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_yaml_max_batch_file_size_mb_wins_over_db(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"max_batch_file_size_mb": 3}, + ) + pc = ProxyConfig() + pc._yaml_general_settings_keys = {"max_batch_file_size_mb"} + await pc._update_general_settings({"max_batch_file_size_mb": 5}) + from litellm.proxy import proxy_server as ps + + assert ps.general_settings.get("max_batch_file_size_mb") == 3 + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_cleared_db_max_batch_file_size_mb_lifts_cap(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"max_batch_file_size_mb": 8}, + ) + pc = ProxyConfig() + await pc._update_general_settings({"max_parallel_requests": 1}) + from litellm.proxy import proxy_server as ps + + assert ps.general_settings.get("max_batch_file_size_mb") is None + + @pytest.mark.asyncio async def test_ProxyConfig__update_general_settings_none_input_noop(): pc = ProxyConfig() From dfeb12649bffb4f3966eb3d32ad2bf413a032404 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 19 Aug 2026 15:45:39 -0700 Subject: [PATCH 61/88] feat(complexity-router): make the reasoning override floor configurable (#37537) The reasoning override's floor was pinned to tier_boundaries.simple_medium, so an operator could not restore the unconditional promotion nor raise the bar independently of the SIMPLE/MEDIUM cut. Setting reasoning_override_min_score was accepted and echoed back by /model/info, because the config model allows extra keys, while routing ignored it. Resolve the floor through one accessor that falls back to simple_medium when the field is unset, so moving that boundary still moves the floor with it, and an explicit 0 is a real floor rather than an absent one. Record the resolved value on the routing decision so a logged row states the floor that applied, which is also what lets the Admin UI stop hardcoding the copy PR #37500 added. --- .../complexity_router/README.md | 2 +- .../complexity_router/complexity_router.py | 17 ++++- .../complexity_router/config.py | 9 +++ litellm/types/utils.py | 2 + .../router_strategy/test_complexity_router.py | 67 +++++++++++++++++++ .../add_model/ClassificationMethodConfig.tsx | 23 +++++-- .../add_model/ComplexityRouterConfig.tsx | 5 ++ .../add_model/HeuristicScoringConfig.test.tsx | 63 +++++++++++++++++ .../add_model/HeuristicScoringConfig.tsx | 59 +++++++++++++++- .../add_model/add_auto_router_tab.test.tsx | 22 ++++++ .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 21 ++++++ .../build_complexity_router_config.ts | 15 ++++- .../add_model/heuristic_scoring_knobs.test.ts | 16 +++++ .../add_model/heuristic_scoring_knobs.ts | 7 ++ ...d_updated_complexity_router_config.test.ts | 36 ++++++++++ .../edit_auto_router_modal.tsx | 7 ++ .../RoutingDecisionCard.test.tsx | 41 +++++++++++- .../LogDetailsDrawer/RoutingDecisionCard.tsx | 17 ++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 ++ 20 files changed, 423 insertions(+), 14 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 259933dbb9e..65f1029bf55 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -165,7 +165,7 @@ response = litellm.completion( ### Reasoning Override -If 2+ reasoning markers are detected in the user message, the request is automatically routed to the REASONING tier regardless of the weighted score. This ensures complex reasoning tasks get the appropriate model. +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. ### System Prompt Handling diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 0573f8acf18..5d6e13c7fc0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1021,10 +1021,10 @@ class ComplexityRouter(CustomLogger): weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) boundaries: Final = self._effective_tier_boundaries() - scored_above_simple: Final = weighted_score >= boundaries["simple_medium"] + clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score() # Reuse match count from _score_keyword_match to avoid scanning twice - if reasoning_match_count >= 2 and scored_above_simple: + if reasoning_match_count >= 2 and clears_override_floor: return ComplexityTier.REASONING, weighted_score, tuple(signals), "reasoning_override" # Map score to tier @@ -1039,6 +1039,18 @@ class ComplexityRouter(CustomLogger): return tier, weighted_score, tuple(signals), "heuristic_scorer" + def _effective_reasoning_override_min_score(self) -> float: + """The score a request must reach before the reasoning-marker override may promote it. + + Unset tracks the SIMPLE/MEDIUM boundary, so moving that boundary moves this floor with it + and the override still cannot rescue a request the mapping would call SIMPLE. An explicit + 0 is a real floor, not an absent one, so the comparison is against None. + """ + configured: Final = self.config.reasoning_override_min_score + if configured is None: + return self._effective_tier_boundaries()["simple_medium"] + return configured + def _effective_tier_boundaries(self) -> StandardLoggingRoutingDecisionTierBoundaries: """The tier boundaries in effect, with the documented defaults filled in. @@ -1095,6 +1107,7 @@ class ComplexityRouter(CustomLogger): if score is not None: decision["score"] = score decision["tier_boundaries"] = self._effective_tier_boundaries() + decision["reasoning_override_min_score"] = self._effective_reasoning_override_min_score() if signals: # Stored as a list because this record is serialized to JSON for the spend # log and read back as an array by the dashboard; a sequence type that only diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 6d43199c948..e82ae991100 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -481,6 +481,15 @@ class ComplexityRouterConfig(BaseModel): ), ) + reasoning_override_min_score: float | None = Field( + default=None, + description=( + "Minimum weighted score a request must reach before 2+ reasoning markers may promote it to the " + "reasoning tier. Unset tracks tier_boundaries.simple_medium, so the override never rescues a " + "request the scorer placed in the cheapest tier; 0 restores the unconditional override" + ), + ) + # Token count thresholds token_thresholds: dict[str, int] = Field( default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 07005d7f9ad..b3e520f8b45 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2836,6 +2836,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_cost: float escalated: bool tier_boundaries: StandardLoggingRoutingDecisionTierBoundaries + reasoning_override_min_score: ReadOnly[float] conversation_continuing: bool savings_baseline_model: str savings_baseline_deployment_id: str @@ -2860,6 +2861,7 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_cost", "escalated", "tier_boundaries", + "reasoning_override_min_score", "conversation_continuing", "savings_baseline_model", "savings_baseline_deployment_id", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index a148da8b675..7586bbf551e 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -285,6 +285,73 @@ class TestReasoningMarkerScoring: assert score == complexity_router.config.tier_boundaries["simple_medium"] assert tier == ComplexityTier.REASONING + def test_explicit_zero_floor_restores_the_unconditional_override(self, mock_router_instance, basic_config): + """0 is a real floor, not an absent one, so the markers alone promote again.""" + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "reasoning_override_min_score": 0.0}, + ) + tier, score, _ = router.classify("hi, step by step, pros and cons") + assert score < router.config.tier_boundaries["simple_medium"] + assert tier == ComplexityTier.REASONING + + def test_floor_defaults_to_simple_medium_and_follows_it(self, mock_router_instance, basic_config): + """Unset tracks simple_medium, so moving that boundary moves the floor with it.""" + prompt = ( + "Give me the pros and cons, step by step, of moving our checkout service " + "to an event-driven architecture." + ) + low = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "tier_boundaries": {"simple_medium": 0.20}}, + ) + high = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "tier_boundaries": {"simple_medium": 0.30}}, + ) + assert low._effective_reasoning_override_min_score() == 0.20 + assert high._effective_reasoning_override_min_score() == 0.30 + assert low.classify(prompt)[0] == ComplexityTier.REASONING + assert high.classify(prompt)[0] != ComplexityTier.REASONING + + def test_explicit_floor_overrides_the_boundary(self, mock_router_instance, basic_config): + """A configured floor decides the override, not simple_medium.""" + prompt = ( + "Give me the pros and cons, step by step, of moving our checkout service " + "to an event-driven architecture." + ) + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "tier_boundaries": {"simple_medium": 0.10}, + "reasoning_override_min_score": 0.90, + }, + ) + tier, score, _ = router.classify(prompt) + assert score > router.config.tier_boundaries["simple_medium"] + assert router._effective_reasoning_override_min_score() == 0.90 + assert tier != ComplexityTier.REASONING + + def test_configured_floor_is_applied_with_greater_or_equal(self, mock_router_instance, basic_config): + """A score landing exactly on the configured floor still promotes.""" + prompt = ( + "Give me the pros and cons, step by step, of moving our checkout service " + "to an event-driven architecture." + ) + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "reasoning_override_min_score": 0.25}, + ) + tier, score, _ = router.classify(prompt) + assert score == 0.25 + assert tier == ComplexityTier.REASONING + def test_system_prompt_reasoning_not_counted(self, complexity_router): """Reasoning markers in system prompt should not count for override.""" user_prompt = "What is 2+2?" diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index da0c50c7bd0..be4573f8a85 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -58,17 +58,32 @@ const scoringExplanation = (value: ComplexityRouterConfigValue): string => { const boundaryRanges = ( shipped: Record | undefined, overrides: Record | undefined, -): { simpleMedium: string; mediumComplex: string; complexReasoning: string } | null => { + reasoningOverrideMinScore: number | undefined, +): { + simpleMedium: string; + mediumComplex: string; + complexReasoning: string; + reasoningOverrideFloor: string; +} | null => { const effective: Record = { ...shipped, ...overrides }; const [low, mid, high] = [effective.simple_medium, effective.medium_complex, effective.complex_reasoning]; if (low === undefined || mid === undefined || high === undefined) return null; - return { simpleMedium: low.toFixed(2), mediumComplex: mid.toFixed(2), complexReasoning: high.toFixed(2) }; + return { + simpleMedium: low.toFixed(2), + mediumComplex: mid.toFixed(2), + complexReasoning: high.toFixed(2), + reasoningOverrideFloor: (reasoningOverrideMinScore ?? low).toFixed(2), + }; }; const HowClassificationWorks: React.FC<{ value: ComplexityRouterConfigValue }> = ({ value }) => { // The shipped boundaries come from the proxy, so this card cannot state ranges the router stopped using. const { data: scorerDefaults, isError } = useComplexityScorerDefaults(); - const ranges = boundaryRanges(scorerDefaults?.tier_boundaries, value.tier_boundaries); + const ranges = boundaryRanges( + scorerDefaults?.tier_boundaries, + value.tier_boundaries, + value.reasoning_override_min_score, + ); return ( @@ -93,7 +108,7 @@ const HowClassificationWorks: React.FC<{ value: ComplexityRouterConfigValue }> =
  • {effectiveTierLabel("REASONING", value.tier_labels)}: Score > {ranges.complexReasoning}{" "} - (or 2+ reasoning markers with a score of at least {ranges.simpleMedium}) + (or 2+ reasoning markers with a score of at least {ranges.reasoningOverrideFloor})
    • )} diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 9f1108639ac..744c74cecb0 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -136,6 +136,11 @@ export interface ComplexityRouterConfigValue { tier_boundaries?: TierBoundaries; token_thresholds?: TokenThresholds; dimension_weights?: DimensionWeights; + /** + * Score floor the reasoning-marker override must clear. Undefined keeps the key out of the payload, so the + * floor tracks tier_boundaries.simple_medium; an explicit 0 is a real floor that promotes on the markers alone. + */ + reasoning_override_min_score?: number; } interface ComplexityRouterConfigProps { diff --git a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx index f1a79ab7c33..9dccd767e49 100644 --- a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx @@ -95,6 +95,57 @@ describe("HeuristicScoringConfig", () => { expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ token_thresholds: undefined })); }); + it("shows the boundary an untouched override floor tracks, rather than a fixed number", async () => { + await render(BASE); + + const field = screen.getByLabelText("Minimum score"); + expect(field).toHaveValue(""); + expect(field).toHaveAttribute("placeholder", SHIPPED_SCORER_DEFAULTS.tier_boundaries.simple_medium.toFixed(2)); + }); + + it("tracks the operator's own Simple to Medium override, not the shipped boundary", async () => { + await render({ ...BASE, tier_boundaries: { simple_medium: 0.42, medium_complex: 0.5, complex_reasoning: 0.7 } }); + + expect(screen.getByLabelText("Minimum score")).toHaveAttribute("placeholder", "0.42"); + }); + + // 0 restores an unconditional override, so it has to reach the config as 0 rather than as "untouched". + it("commits an explicit 0 override floor", async () => { + const onChange = await render(BASE); + fireEvent.change(screen.getByLabelText("Minimum score"), { target: { value: "0" } }); + + expect(onChange.mock.calls.at(-1)?.[0]).toMatchObject({ reasoning_override_min_score: 0 }); + }); + + it("renders a stored 0 as 0 rather than as an untouched field", async () => { + await render({ ...BASE, reasoning_override_min_score: 0 }); + + expect(screen.getByLabelText("Minimum score")).toHaveValue("0"); + }); + + it("counts a set override floor among the overrides", () => { + renderWithProviders( + , + ); + + expect(screen.getByTestId("advanced-scoring-override-count")).toHaveTextContent("1 override"); + }); + + it("clamps the override floor to the score range", async () => { + const onChange = await render(BASE); + fireEvent.change(screen.getByLabelText("Minimum score"), { target: { value: "9" } }); + + expect(onChange.mock.calls.at(-1)?.[0]).toMatchObject({ reasoning_override_min_score: 1 }); + }); + + it("resets the override floor back to tracking the boundary", async () => { + const onChange = await render({ ...BASE, reasoning_override_min_score: 0 }); + + await userEvent.click(screen.getAllByRole("button", { name: "Reset to defaults" }).at(-1)!); + + expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ reasoning_override_min_score: undefined })); + }); + it("flags decreasing boundaries as an error without blocking the save", async () => { const bad = { ...BASE, tier_boundaries: { simple_medium: 0.5, medium_complex: 0.2, complex_reasoning: 0.6 } }; await render(bad); @@ -136,6 +187,18 @@ describe("ClassificationMethodConfig scorer gating", () => { expect(screen.queryByText(/0.15/)).not.toBeInTheDocument(); }); + it("states the configured override floor in the reasoning-marker aside, not the boundary", () => { + renderWithProviders(); + + expect(screen.getByText(/2\+ reasoning markers with a score of at least 0\.00/)).toBeInTheDocument(); + }); + + it("falls back to the Simple to Medium boundary when no override floor is set", () => { + renderWithProviders(); + + expect(screen.getByText(/2\+ reasoning markers with a score of at least 0\.15/)).toBeInTheDocument(); + }); + it("renders a row for every scored dimension", async () => { await render(BASE); diff --git a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.tsx b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.tsx index 1b0a781a33a..540beff8177 100644 --- a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.tsx @@ -23,6 +23,8 @@ interface GroupSpec { labels: Record; } +const OVERRIDE_FLOOR_ID = "reasoning-override-min-score"; + const GROUPS: GroupSpec[] = [ { group: "tier_boundaries", @@ -88,7 +90,12 @@ const HeuristicScoringConfig: React.FC = ({ value, // falls back to the default model, so there is nothing here to configure. const scorerRuns = heuristicScoringRole(value) !== "never"; - const overrides = GROUPS.filter((spec) => value[spec.group] !== undefined).length; + // What an untouched override floor follows: the boundary in effect, override included, not the shipped one. + const trackedFloor: number | undefined = { ...defaults?.tier_boundaries, ...value.tier_boundaries }.simple_medium; + + const overrides = + GROUPS.filter((spec) => value[spec.group] !== undefined).length + + (value.reasoning_override_min_score !== undefined ? 1 : 0); // min/max are inert on a text input, and a plain number input renders Number("0.") as "0" so a decimal // cannot be typed. Hence the local draft plus an explicit clamp here. @@ -102,6 +109,12 @@ const HeuristicScoringConfig: React.FC = ({ value, }); }; + const commitOverrideFloor = (raw: string) => { + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ ...value, reasoning_override_min_score: Math.min(1, Math.max(-1, parsed)) }); + }; + if (!scorerRuns) return null; return ( @@ -215,6 +228,50 @@ const HeuristicScoringConfig: React.FC = ({ value, ); })} + +
    +
    + Reasoning override floor + {value.reasoning_override_min_score !== undefined && ( + + )} +
    +

    + Two or more reasoning markers promote a request to the reasoning tier, but only once its weighted + score reaches this floor.{" "} + {trackedFloor === undefined + ? "Left untouched, it tracks the Simple to Medium boundary." + : `Left untouched, it tracks the Simple to Medium boundary, currently ${trackedFloor.toFixed(2)}.`}{" "} + Set it to 0 to promote on the markers alone. +

    +
    + + { + setDraft({ id: OVERRIDE_FLOOR_ID, raw: event.target.value }); + commitOverrideFloor(event.target.value); + }} + onBlur={() => setDraft(null)} + /> +
    +
    )}
    diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 5643dafc0f5..e3a7482b02d 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -285,6 +285,28 @@ describe("AddAutoRouterTab", () => { }); }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create + // payload is only proven end to end. 0 is the case a truthy check would silently drop. + it("carries a reasoning override floor of 0 through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "override-floor-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(await screen.findByText("Advanced scoring")); + fireEvent.change(await screen.findByLabelText("Minimum score"), { target: { value: "0" } }); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + reasoning_override_min_score: 0, + }); + }); + it("carries session affinity turned on through to the create payload", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); 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 e35b9581d9c..cca40f70833 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 @@ -361,6 +361,7 @@ const AddAutoRouterTab: React.FC = ({ tierBoundaries: complexityRouterConfig.tier_boundaries, tokenThresholds: complexityRouterConfig.token_thresholds, dimensionWeights: complexityRouterConfig.dimension_weights, + reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score, }; const submitRecommendedRouter = async (name: string) => { 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 81d54a7b773..3e33136a3cc 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 @@ -601,6 +601,27 @@ describe("buildComplexityRouterConfig scorer knobs", () => { it("drops them when the classifier falls back to the default model and nothing is scored", () => { expect(buildComplexityRouterConfig(llmWithDefaultFallback)).not.toHaveProperty("tier_boundaries"); }); + + it("omits the reasoning override floor while untouched, so it keeps tracking simple_medium", () => { + expect(buildComplexityRouterConfig(baseParams)).not.toHaveProperty("reasoning_override_min_score"); + }); + + it("emits the reasoning override floor that was set", () => { + const config = buildComplexityRouterConfig({ ...baseParams, reasoningOverrideMinScore: 0.4 }); + expect(config.reasoning_override_min_score).toBe(0.4); + }); + + // 0 is an unconditional override, not an absent knob, so a falsy check here would silently discard it. + it("emits an explicit 0 reasoning override floor", () => { + const config = buildComplexityRouterConfig({ ...baseParams, reasoningOverrideMinScore: 0 }); + expect(config.reasoning_override_min_score).toBe(0); + }); + + it("drops the reasoning override floor when nothing is scored", () => { + expect(buildComplexityRouterConfig({ ...llmWithDefaultFallback, reasoningOverrideMinScore: 0 })).not.toHaveProperty( + "reasoning_override_min_score", + ); + }); }); describe("plan-mode minimum tier", () => { 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 0ab2db8c16f..bddf8321ad2 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 @@ -46,6 +46,7 @@ interface ScorerKnobInputs { tierBoundaries: TierBoundaries | undefined; tokenThresholds: TokenThresholds | undefined; dimensionWeights: DimensionWeights | undefined; + reasoningOverrideMinScore: number | undefined; } /** @@ -59,6 +60,7 @@ const scorerKnobPayload = ({ tierBoundaries, tokenThresholds, dimensionWeights, + reasoningOverrideMinScore, }: ScorerKnobInputs) => heuristicScoringRoleFor(classifierType, classifierFallback) === "never" ? {} @@ -66,6 +68,7 @@ const scorerKnobPayload = ({ ...(tierBoundaries && { tier_boundaries: tierBoundaries }), ...(tokenThresholds && { token_thresholds: tokenThresholds }), ...(dimensionWeights && { dimension_weights: dimensionWeights }), + ...(reasoningOverrideMinScore !== undefined && { reasoning_override_min_score: reasoningOverrideMinScore }), }; export interface BuildComplexityRouterConfigParams { @@ -95,6 +98,7 @@ export interface BuildComplexityRouterConfigParams { tierBoundaries?: TierBoundaries; tokenThresholds?: TokenThresholds; dimensionWeights?: DimensionWeights; + reasoningOverrideMinScore?: number; } export interface ComplexityRouterConfigPayload { @@ -124,6 +128,7 @@ export interface ComplexityRouterConfigPayload { tier_boundaries?: TierBoundaries; token_thresholds?: TokenThresholds; dimension_weights?: DimensionWeights; + reasoning_override_min_score?: number; } const TIER_KEYS: Array = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]; @@ -230,11 +235,19 @@ export const buildComplexityRouterConfig = ({ tierBoundaries, tokenThresholds, dimensionWeights, + reasoningOverrideMinScore, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const cleanedEscalationKeywords = escalationKeywords.map((keyword) => keyword.trim()).filter(Boolean); const cleanedKeywordTierRules = serializeKeywordTierRules(keywordTierRules); const cleanedTierLabels = serializeTierLabels(tierLabels); - const scorerInputs = { classifierType, classifierFallback, tierBoundaries, tokenThresholds, dimensionWeights }; + const scorerInputs = { + classifierType, + classifierFallback, + tierBoundaries, + tokenThresholds, + dimensionWeights, + reasoningOverrideMinScore, + }; const scorerKnobs = scorerKnobPayload(scorerInputs); return { diff --git a/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.test.ts b/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.test.ts index 4f6c92e5a35..ee1d819872b 100644 --- a/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.test.ts @@ -4,6 +4,7 @@ import { heuristicScoringRoleFor } from "./ComplexityRouterConfig"; import { dimensionLabel, hydrateDimensionWeights, + hydrateReasoningOverrideMinScore, hydrateTierBoundaries, hydrateTokenThresholds, weightTotal, @@ -46,6 +47,21 @@ describe("hydrating the scorer knobs", () => { expect(weightTotal({ a: 0.1, b: 0.2 })).toBe(0.3); }); + it.each([[undefined], [null], ["0.15"], [Number.NaN], [Number.POSITIVE_INFINITY], [{ value: 0.15 }]])( + "hydrates the reasoning override floor %s to undefined", + (raw) => { + expect(hydrateReasoningOverrideMinScore(raw)).toBeUndefined(); + }, + ); + + // A stored 0 is an unconditional override, so hydrating it to undefined would silently retune the router + // back to tracking simple_medium on the next save. + it("hydrates a stored reasoning override floor, zero and negatives included", () => { + expect(hydrateReasoningOverrideMinScore(0)).toBe(0); + expect(hydrateReasoningOverrideMinScore(-0.3)).toBe(-0.3); + expect(hydrateReasoningOverrideMinScore(0.42)).toBe(0.42); + }); + it("falls back to the raw key when a dimension has no label yet", () => { expect(dimensionLabel("codePresence")).toBe("Code presence"); expect(dimensionLabel("somethingNew")).toBe("somethingNew"); diff --git a/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.ts b/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.ts index 46e145554ac..1f44a479c60 100644 --- a/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.ts +++ b/ui/litellm-dashboard/src/components/add_model/heuristic_scoring_knobs.ts @@ -44,5 +44,12 @@ export const hydrateTokenThresholds = (raw: unknown): TokenThresholds | undefine export const hydrateDimensionWeights = (raw: unknown): DimensionWeights | undefined => hydrateNumericMap(raw); +/** + * The scalar counterpart of hydrateNumericMap: absent hydrates to undefined so an untouched save keeps the + * floor tracking tier_boundaries.simple_medium, while a stored 0 hydrates to 0, which is a real floor. + */ +export const hydrateReasoningOverrideMinScore = (raw: unknown): number | undefined => + typeof raw === "number" && Number.isFinite(raw) ? raw : undefined; + export const weightTotal = (weights: DimensionWeights): number => Math.round(Object.values(weights).reduce((total, weight) => total + weight, 0) * 100) / 100; 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 c62955b6683..0ae782dc5fc 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 @@ -313,6 +313,42 @@ describe("buildUpdatedComplexityRouterConfig scorer knobs", () => { it("never invents knobs for a router that never had them", () => { expect(buildUpdatedComplexityRouterConfig(STORED, FORM_VALUE)).not.toHaveProperty("tier_boundaries"); }); + + // 0 is an unconditional reasoning override. Treating it as unset here would quietly retune the router + // back to tracking simple_medium on the next save. + it("round-trips a stored reasoning override floor of 0", () => { + const result = buildUpdatedComplexityRouterConfig( + { ...STORED, reasoning_override_min_score: 0 }, + { ...FORM_VALUE, reasoning_override_min_score: 0 }, + ); + expect(result.reasoning_override_min_score).toBe(0); + }); + + it("writes a newly set reasoning override floor over the stored one", () => { + const result = buildUpdatedComplexityRouterConfig( + { ...STORED, reasoning_override_min_score: 0 }, + { ...FORM_VALUE, reasoning_override_min_score: 0.5 }, + ); + expect(result.reasoning_override_min_score).toBe(0.5); + }); + + it("drops a stored reasoning override floor when the operator resets it", () => { + const result = buildUpdatedComplexityRouterConfig({ ...STORED, reasoning_override_min_score: 0.5 }, FORM_VALUE); + + expect(result).not.toHaveProperty("reasoning_override_min_score"); + expect(result.some_future_backend_key).toEqual({ nested: true }); + }); + + it("drops the reasoning override floor on a router whose scorer never runs", () => { + const neverScores = { + ...FORM_VALUE, + reasoning_override_min_score: 0, + classifier_type: "llm" as const, + classifier_fallback: "default_model" as const, + }; + const result = buildUpdatedComplexityRouterConfig({ ...STORED, reasoning_override_min_score: 0 }, neverScores); + expect(result).not.toHaveProperty("reasoning_override_min_score"); + }); }); describe("buildUpdatedComplexityRouterConfig plan-mode minimum tier", () => { 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 54f992b08f4..012624454d3 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 @@ -30,6 +30,7 @@ import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; import { hydrateKeywordTierRules, serializeKeywordTierRules } from "../add_model/complexity_router_keywords"; import { hydrateDimensionWeights, + hydrateReasoningOverrideMinScore, hydrateTierBoundaries, hydrateTokenThresholds, } from "../add_model/heuristic_scoring_knobs"; @@ -84,6 +85,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "tier_boundaries", "token_thresholds", "dimension_weights", + "reasoning_override_min_score", ]); // Managed only when the caller passes the corresponding state. A caller that does not render @@ -200,6 +202,10 @@ export const buildUpdatedComplexityRouterConfig = ( ...(scorerRuns && value.tier_boundaries !== undefined && { tier_boundaries: value.tier_boundaries }), ...(scorerRuns && value.token_thresholds !== undefined && { token_thresholds: value.token_thresholds }), ...(scorerRuns && value.dimension_weights !== undefined && { dimension_weights: value.dimension_weights }), + ...(scorerRuns && + value.reasoning_override_min_score !== undefined && { + reasoning_override_min_score: value.reasoning_override_min_score, + }), }; }; @@ -367,6 +373,7 @@ const EditAutoRouterModal: React.FC = ({ tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), + reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), session_affinity: typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index ed99e414706..b6803f23f50 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -60,7 +60,9 @@ describe("RoutingDecisionCard", () => { />, ); expect( - screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score above the lowest tier)"), + screen.getByText( + "Heuristic, REASONING override (2 or more reasoning markers, score of at least the Simple to Medium boundary)", + ), ).toBeInTheDocument(); expect(screen.getByText("0.20")).toBeInTheDocument(); // The score did not decide this tier, so NO band explanation may render at all. @@ -191,7 +193,9 @@ describe("RoutingDecisionCard", () => { render(); expect(screen.queryByText(/SIMPLE|MEDIUM|COMPLEX|at or above/)).not.toBeInTheDocument(); expect( - screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score above the lowest tier)"), + screen.getByText( + "Heuristic, REASONING override (2 or more reasoning markers, score of at least the Simple to Medium boundary)", + ), ).toBeInTheDocument(); }); @@ -217,10 +221,41 @@ describe("RoutingDecisionCard", () => { , ); expect( - screen.getByText("Heuristic, Deep override (2 or more reasoning markers, score above the lowest tier)"), + screen.getByText( + "Heuristic, Deep override (2 or more reasoning markers, score of at least the Simple to Medium boundary)", + ), ).toBeInTheDocument(); }); + it("states the floor the override actually cleared", () => { + render( + , + ); + expect( + screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score of at least 0.05)"), + ).toBeInTheDocument(); + }); + + // A floor of 0 is an unconditional override, so a falsy check here would print the "before this change" + // wording on a row that recorded a real floor. + it("states a recorded floor of 0 rather than treating it as unrecorded", () => { + render( + , + ); + expect( + screen.getByText("Heuristic, REASONING override (2 or more reasoning markers, score of at least 0)"), + ).toBeInTheDocument(); + }); + + it("never prints undefined on a row logged before the floor was recorded", () => { + render(); + expect(screen.queryByText(/undefined/)).not.toBeInTheDocument(); + }); + it("falls back to the raw cause for a value this build does not know", () => { render(); expect(screen.getByText("some_future_cause")).toBeInTheDocument(); 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 cb680668849..eaee4d3ce23 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -26,6 +26,7 @@ export interface RoutingDecision { classifier_model?: string; escalated?: boolean; tier_boundaries?: RoutingDecisionTierBoundaries; + reasoning_override_min_score?: number; } const ROUTER_TYPE_LABELS: Record = { @@ -65,14 +66,26 @@ function describePlanModeFloor(matchedKeyword: string | undefined): string { return "Plan-mode floor"; } +/** Rows logged before the floor was recorded name what it tracked back then instead of a number. */ +function describeReasoningOverride(tierLabel: string | undefined, floor: number | undefined): string { + const stated = floor === undefined ? "the Simple to Medium boundary" : String(floor); + return `Heuristic, ${tierLabel ?? "REASONING"} override (2 or more reasoning markers, score of at least ${stated})`; +} + function describeCause(decision: RoutingDecision): string { - const { cause, classifier_model: classifierModel, matched_keyword: matchedKeyword, tier_label: tierLabel } = decision; + const { + cause, + classifier_model: classifierModel, + matched_keyword: matchedKeyword, + tier_label: tierLabel, + reasoning_override_min_score: overrideFloor, + } = decision; switch (cause) { case "heuristic_scorer": return "Heuristic scorer"; case "reasoning_override": - return `Heuristic, ${tierLabel ?? "REASONING"} override (2 or more reasoning markers, score above the lowest tier)`; + return describeReasoningOverride(tierLabel, overrideFloor); case "llm_classifier": return classifierModel ? `LLM classifier (${classifierModel})` : "LLM classifier"; case "literal_keyword_match": diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4c63586d4d6..146ee1a723a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32480,6 +32480,11 @@ export interface components { * @description Keywords indicating reasoning-required content */ reasoning_keywords?: string[] | null; + /** + * Reasoning Override Min Score + * @description Minimum weighted score a request must reach before 2+ reasoning markers may promote it to the reasoning tier. Unset tracks tier_boundaries.simple_medium, so the override never rescues a request the scorer placed in the cheapest tier; 0 restores the unconditional override + */ + reasoning_override_min_score?: number | null; /** * Reminder Markers * @description Override the delimiter pairs used to recognize and strip harness-injected reminder blocks before classification. A harness that wraps injected context differently per agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than adds to, the built-in default of ('', ''), so a harness that also emits that pair lists it too. Matching is case-insensitive. @@ -33502,6 +33507,8 @@ export interface components { escalation_keyword?: string; /** Matched Keyword */ matched_keyword?: string; + /** Reasoning Override Min Score */ + reasoning_override_min_score?: number; /** Request Type */ request_type?: string; /** Routed Model */ From 5d6033f8b4b65a8dd1c67a75695deb20112f7f72 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 15:46:55 -0700 Subject: [PATCH 62/88] refactor(ui): migrate the remaining dashboard pages off antd (#37524) * refactor(ui): migrate the remaining dashboard pages off antd Converts the teams, usage, guardrails, vector stores, cost tracking, agents, policies, login and onboarding screens onto the shadcn primitives, including the team info tab shell and the virtual keys hover cards. * fix(ui): close out the antd migration's failing type checks and tests Alert and Badge were missing the success and info variants their call sites already used. Combobox dropped disabled because Base UI merges the primitive's own props over the render child, so the flag never reached the input, and the guardrails status filter had no accessible name, which left two comboboxes indistinguishable to the tests. The remaining test updates swap antd's title-based queries for the roles the shadcn controls expose. --- .../agents/_components/agent_info.tsx | 106 +- .../add_provider_form.integration.test.tsx | 51 +- ...ost_tracking_settings.integration.test.tsx | 6 +- .../cost_tracking_settings.test.tsx | 4 +- .../_components/GuardrailConfig.test.tsx | 14 +- .../_components/GuardrailConfig.tsx | 127 +- .../TeamGuardrailsTab.integration.test.tsx | 6 +- .../_components/TeamGuardrailsTab.tsx | 1 + .../components/ModelSelector.test.tsx | 27 +- .../compareUI/components/ModelSelector.tsx | 28 +- .../_components/policy_test_panel.tsx | 8 +- .../_components/CreateVectorStore.tsx | 171 +- .../_components/S3VectorsConfig.tsx | 34 +- .../src/app/login/LoginPage.tsx | 11 +- .../src/app/onboarding/OnboardingFormBody.tsx | 2 +- .../LoggingSettings/LoggingSettings.test.tsx | 4 +- .../LoggingSettings/LoggingSettings.tsx | 46 +- .../MCPSemanticFilterSettings.tsx | 273 +-- .../PluginSettings.integration.test.tsx | 2 +- .../SSOSettings/RoleMappings.tsx | 115 +- .../EntityUsage/TopKeyView.test.tsx | 12 +- .../components/EntityUsage/TopKeyView.tsx | 31 +- .../VirtualKeysPage/keyTableColumns.tsx | 38 +- .../agent_management/AgentSelector.test.tsx | 72 +- .../agent_management/AgentSelector.tsx | 58 +- .../create_key_button.integration.test.tsx | 4 +- .../src/components/shared/Alert.tsx | 1 + .../components/shared/SearchSelect.test.tsx | 16 + .../components/team/LoggingSettings.test.tsx | 4 +- .../src/components/team/LoggingSettings.tsx | 128 +- .../src/components/team/TeamInfo.tsx | 2000 ++++++++--------- .../src/components/ui/badge.tsx | 1 + .../src/components/ui/combobox.tsx | 2 +- 33 files changed, 1733 insertions(+), 1670 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index d37065317c2..c507d56a63b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect, useMemo } from "react"; -import { Spin, Descriptions } from "antd"; +import { cx } from "@/lib/cva.config"; import { FormProvider, useForm, useWatch } from "react-hook-form"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; @@ -40,6 +40,26 @@ interface AgentInfoViewProps { isAdmin: boolean; } +const DetailList: React.FC<{ children: React.ReactNode; className?: string }> = ({ children, className }) => ( +
    + {children} +
    +); + +const DetailItem: React.FC<{ label: React.ReactNode; children: React.ReactNode }> = ({ label, children }) => ( + <> +
    + {label} +
    +
    {children}
    + +); + const AgentInfoView: React.FC = ({ agentId, onClose, accessToken, isAdmin }) => { const [agent, setAgent] = useState(null); const [selectedKey, setSelectedKey] = useState(null); @@ -195,7 +215,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT return (
    - +
    ); @@ -276,51 +296,41 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT
    {/* Overview Panel */} - - {agent.agent_id} - {agent.agent_name} - {agent.agent_card_params?.name || "-"} - {agent.agent_card_params?.description || "-"} - {agent.agent_card_params?.url || "-"} - {agent.agent_card_params?.version || "-"} - - {agent.agent_card_params?.protocolVersion || "-"} - - + + {agent.agent_id} + {agent.agent_name} + {agent.agent_card_params?.name || "-"} + {agent.agent_card_params?.description || "-"} + {agent.agent_card_params?.url || "-"} + {agent.agent_card_params?.version || "-"} + {agent.agent_card_params?.protocolVersion || "-"} + {agent.agent_card_params?.capabilities?.streaming ? "Yes" : "No"} - + {agent.agent_card_params?.capabilities?.pushNotifications && ( - Yes + Yes )} {agent.agent_card_params?.capabilities?.stateTransitionHistory && ( - Yes - )} - - {agent.agent_card_params?.skills?.length || 0} configured - - {agent.litellm_params?.model && ( - {agent.litellm_params.model} + Yes )} + {agent.agent_card_params?.skills?.length || 0} configured + {agent.litellm_params?.model && {agent.litellm_params.model}} {agent.litellm_params?.make_public !== undefined && ( - - {agent.litellm_params.make_public ? "Yes" : "No"} - + {agent.litellm_params.make_public ? "Yes" : "No"} )} {agent.agent_card_params?.iconUrl && ( - {agent.agent_card_params.iconUrl} + {agent.agent_card_params.iconUrl} )} {agent.agent_card_params?.documentationUrl && ( - - {agent.agent_card_params.documentationUrl} - + {agent.agent_card_params.documentationUrl} )} - {agent.tpm_limit ?? "Unlimited"} - {agent.rpm_limit ?? "Unlimited"} - {agent.session_tpm_limit ?? "Unlimited"} - {agent.session_rpm_limit ?? "Unlimited"} - {formatDate(agent.created_at)} - {formatDate(agent.updated_at)} - + {agent.tpm_limit ?? "Unlimited"} + {agent.rpm_limit ?? "Unlimited"} + {agent.session_tpm_limit ?? "Unlimited"} + {agent.session_rpm_limit ?? "Unlimited"} + {formatDate(agent.created_at)} + {formatDate(agent.updated_at)} + @@ -331,21 +341,19 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (

    MCP Tool Permissions

    - + {agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && ( - - {agent.object_permission.mcp_servers.join(", ")} - + {agent.object_permission.mcp_servers.join(", ")} )} {agent.object_permission.mcp_access_groups && agent.object_permission.mcp_access_groups.length > 0 && ( - + {agent.object_permission.mcp_access_groups.join(", ")} - + )} {agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && ( - +
    {Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
    @@ -354,9 +362,9 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT
    ))}
    -     + )} -     +
    )} @@ -365,9 +373,9 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.agent_card_params?.skills && agent.agent_card_params.skills.length > 0 && (

    Skills

    - + {agent.agent_card_params.skills.map((skill: any, index: number) => ( - +
    ID: {skill.id} @@ -385,9 +393,9 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT
    )}
    -     + ))} -     +
    )}     diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx index 83875cfc676..4680a4d504a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.integration.test.tsx @@ -1,5 +1,3 @@ -// eslint-disable-next-line no-restricted-imports -- the parent cost_tracking_settings still owns this antd Form, and the point of this test is that AddProviderForm registers nothing in it -import { Form } from "antd"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -10,42 +8,31 @@ import { DiscountConfig } from "./types"; const onAddProvider = vi.fn(); const onParentFinish = vi.fn(); -const readParentStore = vi.fn(); -const ParentOwnedForm = () => { - const [form] = Form.useForm(); - return ( - - { - readParentStore(form.getFieldsValue()); - onAddProvider(); - }} - /> - - ); -}; +const ParentOwnedForm = () => ( +
    { + event.preventDefault(); + onParentFinish(); + }} + className="space-y-6" + > + + +); -describe("AddProviderForm inside the antd form its parent owns", () => { +describe("AddProviderForm inside the form its parent owns", () => { beforeEach(() => { vi.clearAllMocks(); }); - it("registers no field in the parent FormInstance, so the parent's resetFields is a no-op", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - await user.click(screen.getByRole("button", { name: /add provider discount/i })); - - expect(readParentStore).toHaveBeenCalledTimes(1); - expect(readParentStore.mock.calls[0][0]).toEqual({}); - }); - it("drives both the onAddProvider prop and the parent form submit from one click", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx index f811b55bc86..d4fbb7e153e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx @@ -72,7 +72,7 @@ describe("CostTrackingSettings submit paths", () => { const header = screen.getByText("Provider Discounts").closest("button"); if (header) await user.click(header); await user.click(await screen.findByRole("button", { name: /add provider discount/i })); - await screen.findByText("Add Provider Discount", { selector: "h2" }); + await screen.findByRole("dialog", { name: "Add Provider Discount" }); }; const submitDiscount = () => @@ -101,7 +101,7 @@ describe("CostTrackingSettings submit paths", () => { const header = screen.getByText("Fee/Price Margin").closest("button"); if (header) await user.click(header); await user.click(await screen.findByRole("button", { name: /add provider margin/i })); - await screen.findByText("Add Provider Margin", { selector: "h2" }); + await screen.findByRole("dialog", { name: "Add Provider Margin" }); await user.click(screen.getAllByRole("combobox")[0]); await user.click((await screen.findAllByRole("option"))[0]); @@ -136,7 +136,7 @@ describe("CostTrackingSettings submit paths", () => { const header = screen.getByText("Fee/Price Margin").closest("button"); if (header) await user.click(header); await user.click(await screen.findByRole("button", { name: /add provider margin/i })); - await screen.findByText("Add Provider Margin", { selector: "h2" }); + await screen.findByRole("dialog", { name: "Add Provider Margin" }); await user.click(screen.getAllByRole("combobox")[0]); await user.click((await screen.findAllByRole("option"))[0]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx index 9cdc8509fbf..a8609aef629 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx @@ -136,7 +136,7 @@ describe("CostTrackingSettings", () => { const addButton = await screen.findByRole("button", { name: /add provider discount/i }); await user.click(addButton); - expect(await screen.findByText("Add Provider Discount", { selector: "h2" })).toBeInTheDocument(); + expect(await screen.findByRole("dialog", { name: "Add Provider Discount" })).toBeInTheDocument(); }); }); @@ -153,7 +153,7 @@ describe("CostTrackingSettings", () => { const addButton = await screen.findByRole("button", { name: /add provider margin/i }); await user.click(addButton); - expect(await screen.findByText("Add Provider Margin", { selector: "h2" })).toBeInTheDocument(); + expect(await screen.findByRole("dialog", { name: "Add Provider Margin" })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.test.tsx index 78526251049..60bf235040f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.test.tsx @@ -45,19 +45,7 @@ describe("GuardrailConfig", () => { it("should show custom code textarea when custom code override is toggled on", async () => { const user = userEvent.setup(); render(); - // Walk up from "Custom Code Override" heading to find the enclosing section, - // then locate the switch within it - const heading = screen.getByText("Custom Code Override"); - let container = heading.parentElement; - let customCodeSwitch: Element | null = null; - while (container && !customCodeSwitch) { - customCodeSwitch = container.querySelector('[role="switch"]'); - container = container.parentElement; - } - if (!customCodeSwitch) { - throw new Error("Could not find the Custom Code Override switch via DOM traversal"); - } - await user.click(customCodeSwitch); + await user.click(screen.getByRole("switch", { name: "Custom Code Override" })); expect(screen.getByPlaceholderText(/async def evaluate/)).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.tsx index 667c52087f1..619ccd8d974 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailConfig.tsx @@ -5,9 +5,13 @@ import { RollbackOutlined, SaveOutlined, } from "@ant-design/icons"; -import { Input, Select, Switch } from "antd"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Button } from "@/components/ui/button"; -import React, { useState } from "react"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; +import { Textarea } from "@/components/ui/textarea"; +import React, { useId, useState } from "react"; interface GuardrailConfigProps { guardrailName: string; @@ -27,6 +31,28 @@ const versions = [ { id: "v1", label: "v1", date: "2026-01-28", author: "admin@company.com", changes: "Initial configuration" }, ]; +const ACTION_ITEMS = [ + { value: "block", label: "Block Request" }, + { value: "flag", label: "Flag for Review" }, + { value: "log", label: "Log Only" }, + { value: "fallback", label: "Use Fallback Response" }, +]; + +const PROVIDER_ITEMS = [ + { value: "bedrock", label: "AWS Bedrock Guardrails" }, + { value: "google", label: "Google Cloud AI Safety" }, + { value: "litellm", label: "LiteLLM Built-in" }, + { value: "custom", label: "Custom Code" }, +]; + +const GUARDRAIL_TYPE_ITEMS = [ + { value: "Content Safety", label: "Content Safety" }, + { value: "PII", label: "PII Detection" }, + { value: "Topic", label: "Topic Restriction" }, + { value: "prompt_injection", label: "Prompt Injection" }, + { value: "custom", label: "Custom" }, +]; + export function GuardrailConfig({ guardrailName, guardrailType, provider }: GuardrailConfigProps) { const [action, setAction] = useState("block"); const [enabled, setEnabled] = useState(true); @@ -35,6 +61,7 @@ export function GuardrailConfig({ guardrailName, guardrailType, provider }: Guar const [rerunStatus, setRerunStatus] = useState<"idle" | "running" | "success" | "error">("idle"); const [version, setVersion] = useState("v3"); const [showVersionHistory, setShowVersionHistory] = useState(false); + const enabledToggleId = useId(); const handleRerun = () => { setRerunStatus("running"); @@ -52,11 +79,21 @@ export function GuardrailConfig({ guardrailName, guardrailType, provider }: Guar
    Version: @@ -109,45 +146,53 @@ export function GuardrailConfig({ guardrailName, guardrailType, provider }: Guar
    - + + + + + {PROVIDER_ITEMS.map((item) => ( + + {item.label} + + ))} + +
    - + + + + + {GUARDRAIL_TYPE_ITEMS.map((item) => ( + + {item.label} + + ))} + +
    @@ -156,8 +201,10 @@ export function GuardrailConfig({ guardrailName, guardrailType, provider }: Guar
    - - Guardrail enabled in production + +
    @@ -172,11 +219,11 @@ export function GuardrailConfig({ guardrailName, guardrailType, provider }: Guar

    Replace the built-in guardrail with custom evaluation code

    - +
    {useCustomCode && ( - setCustomCode(e.target.value)} placeholder={`async def evaluate(input_text: str, context: dict) -> dict: diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.integration.test.tsx index 5c34eca9804..ef1746310ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.integration.test.tsx @@ -136,7 +136,7 @@ describe("TeamGuardrailsTab submit payload", () => { await openSubmitModal(user); await fillRequiredFields(user, "https://guard.example.com/v1/check"); - await user.click(screen.getAllByRole("combobox")[1]); + await user.click(screen.getByRole("combobox", { name: "Mode" })); const options = await screen.findAllByText("During Call"); await user.click(options[options.length - 1]); await submit(user); @@ -149,13 +149,13 @@ describe("TeamGuardrailsTab submit payload", () => { const user = userEvent.setup(); await openSubmitModal(user); - const mode = screen.getAllByRole("combobox")[1]; + const mode = screen.getByRole("combobox", { name: "Mode" }); expect(mode).toHaveTextContent("Pre Call"); await user.click(mode); await user.click(await screen.findByRole("option", { name: "During Call" })); - expect(screen.getAllByRole("combobox")[1]).toHaveTextContent("During Call"); + expect(screen.getByRole("combobox", { name: "Mode" })).toHaveTextContent("During Call"); }); it("blocks an empty submit and reports every required field", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index fc77562e932..4edef302f13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -1001,6 +1001,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { /> + emptyText="No models found" + allowClear={false} + className="rounded-md" + /> {isAddingCustom && ( = ({ accessToken }) => { {hasSearched && result && (
    {result.matched_policies.length === 0 ? ( - +
    +
    ) : ( <>
    diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 271ef31a68e..eae4a3fd799 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -1,9 +1,7 @@ -import React, { useState } from "react"; -import { Upload } from "antd"; +import React, { useId, useState } from "react"; import { toast } from "@/lib/toast"; -import { InboxOutlined } from "@ant-design/icons"; -import type { UploadProps } from "antd"; -import { CircleCheck, CircleHelp, X } from "lucide-react"; +import { CircleCheck, CircleHelp, Inbox, X } from "lucide-react"; +import { v4 as uuidv4 } from "uuid"; import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { ragIngestCall } from "@/components/networking"; import { DocumentUpload, RAGIngestResponse } from "@/components/vector_store_management/types"; @@ -26,7 +24,17 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import S3VectorsConfig from "./S3VectorsConfig"; -const { Dragger } = Upload; +const ACCEPTED_DOCUMENT_EXTENSIONS = ".pdf,.txt,.docx,.md,.doc"; + +const ACCEPTED_DOCUMENT_TYPES = [ + "application/pdf", + "text/plain", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "text/markdown", +]; + +const MAX_DOCUMENT_BYTES = 50 * 1024 * 1024; const RAG_INGEST_UNSUPPORTED_PROVIDERS = new Set(["valkey"]); @@ -39,6 +47,36 @@ const providerItems = Object.entries(VectorStoreProviders) const asText = (value: unknown): string => (typeof value === "string" ? value : ""); +const IngestSuccessAlert: React.FC<{ ingestResults: RAGIngestResponse[] }> = ({ ingestResults }) => { + const [dismissed, setDismissed] = useState(false); + + if (dismissed) { + return null; + } + + return ( + + + Vector Store Created Successfully + +
    +

    + Vector Store ID: {ingestResults[0]?.vector_store_id} +

    +

    + Documents Ingested: {ingestResults.length} +

    +
    +     + + + +     + ); +}; + const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> {label} @@ -62,53 +100,33 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu const [vectorStoreDescription, setVectorStoreDescription] = useState(""); const [ingestResults, setIngestResults] = useState([]); const [providerParams, setProviderParams] = useState>({}); + const documentsInputId = useId(); - const uploadProps: UploadProps = { - name: "file", - multiple: true, - accept: ".pdf,.txt,.docx,.md,.doc", - beforeUpload: (file) => { - const isValidType = [ - "application/pdf", - "text/plain", - "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - "application/msword", - "text/markdown", - ].includes(file.type); + const isSupportedDocument = (file: File): boolean => { + if (!ACCEPTED_DOCUMENT_TYPES.includes(file.type)) { + toast.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`); + return false; + } + if (file.size >= MAX_DOCUMENT_BYTES) { + toast.error(`${file.name} must be smaller than 50MB!`); + return false; + } + return true; + }; - if (!isValidType) { - toast.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`); - return Upload.LIST_IGNORE; - } + const handleAddDocuments = (files: readonly File[]) => { + const accepted: DocumentUpload[] = files.filter(isSupportedDocument).map((file) => ({ + uid: uuidv4(), + name: file.name, + status: "done", + size: file.size, + type: file.type, + originFileObj: file, + })); - const isLt50M = file.size / 1024 / 1024 < 50; - if (!isLt50M) { - toast.error(`${file.name} must be smaller than 50MB!`); - return Upload.LIST_IGNORE; - } - - const newDoc: DocumentUpload = { - uid: file.uid, - name: file.name, - status: "done", - size: file.size, - type: file.type, - originFileObj: file, - }; - - setDocuments((prev) => [...prev, newDoc]); - return false; // Prevent auto upload - }, - onRemove: (file) => { - setDocuments((prev) => prev.filter((doc) => doc.uid !== file.uid)); - }, - fileList: documents.map((doc) => ({ - uid: doc.uid, - name: doc.name, - status: doc.status, - size: doc.size, - })), - showUploadList: false, // We'll use our custom table + if (accepted.length > 0) { + setDocuments((prev) => [...prev, ...accepted]); + } }; const handleRemoveDocument = (uid: string) => { @@ -235,15 +253,32 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu Upload one or more documents (PDF, TXT, DOCX, MD). Maximum file size: 50MB per file.

    - -

    - -

    -

    Click or drag files to this area to upload

    -

    +

    -     + + { + handleAddDocuments(Array.from(event.target.files ?? [])); + event.target.value = ""; + }} + /> + @@ -359,25 +394,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu {/* Success Message */} - {ingestResults.length > 0 && ( - - - Vector Store Created Successfully - -

    - Vector Store ID: {ingestResults[0]?.vector_store_id} -

    -

    - Documents Ingested: {ingestResults.length} -

    -     - - - -     - )} + {ingestResults.length > 0 && }
    ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx index cdc63e8b7ec..7243291d740 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/S3VectorsConfig.tsx @@ -76,22 +76,24 @@ const S3VectorsConfig: React.FC = ({ accessToken, provider AWS S3 Vectors Setup -

    AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

    -
      -
    • Vector buckets and indexes will be automatically created if they don't exist
      • -
    • Vector dimensions are auto-detected from your selected embedding model
      • -
    • Ensure your AWS credentials have permissions for S3 Vectors operations
      • -
    • - Learn more:{" "} - - AWS S3 Vectors Documentation - -
      • -
    +
    +

    AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

    +
      +
    • Vector buckets and indexes will be automatically created if they don't exist
      • +
    • Vector dimensions are auto-detected from your selected embedding model
      • +
    • Ensure your AWS credentials have permissions for S3 Vectors operations
      • +
    • + Learn more:{" "} + + AWS S3 Vectors Documentation + +
      • +
    +
    diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index d303a3da87f..3a1d120c8a5 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -32,8 +32,11 @@ const loginSchema = z.object({ type LoginFormValues = z.infer; function SsoEnabledNotice() { - const [isDismissed, setIsDismissed] = useState(false); - if (isDismissed) return null; + const [dismissed, setDismissed] = useState(false); + + if (dismissed) { + return null; + } return ( @@ -45,8 +48,8 @@ function SsoEnabledNotice() { environment configuration. - diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx index 0a03cc6a591..77fae186c90 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx @@ -1,5 +1,5 @@ -import { CircleAlert, Info } from "lucide-react"; import React from "react"; +import { CircleAlert, Info } from "lucide-react"; import { z } from "zod/v4"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { PasswordInput } from "@/components/shared/PasswordInput"; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx index 46aaa2b4f98..94a6d389db2 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx @@ -500,9 +500,7 @@ describe("LoggingSettings", () => { expect(screen.queryByPlaceholderText("e.g., 7d, 30d")).not.toBeInTheDocument(); expect(screen.queryByRole("button", { name: "Save Settings" })).not.toBeInTheDocument(); - // eslint-disable-next-line local/no-antd-class-selectors -- antd Skeleton exposes no role, label or aria-busy to query the loading affordance by - const skeletons = document.querySelectorAll(".ant-skeleton"); - expect(skeletons.length).toBeGreaterThan(0); + expect(document.querySelectorAll('[data-slot="skeleton"]').length).toBeGreaterThan(0); }); it("should report an error and not claim success when clearing a field fails", async () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx index f67406a4204..20dfb3d96fb 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx @@ -21,7 +21,8 @@ import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ClockCircleOutlined } from "@ant-design/icons"; -import { Card, Skeleton, Space, Typography } from "antd"; +import { Skeleton } from "@/components/ui/skeleton"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { CircleHelp } from "lucide-react"; import React, { useCallback, useMemo } from "react"; import { useForm } from "react-hook-form"; @@ -317,23 +318,34 @@ const LoggingSettings: React.FC = () => { }; return ( - - - - Proxy-wide settings that control how request and response data are written to spend logs. - + + + Logging Settings + + +
    +

    + Proxy-wide settings that control how request and response data are written to spend logs. +

    - {isLoadingConfig ? ( - - ) : ( - - )} - + {isLoadingConfig ? ( +
    + + + + + +
    + ) : ( + + )} +
    +     ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx index ba0b3a5037a..f7ec5238a9d 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx @@ -3,7 +3,7 @@ import { useMCPSemanticFilterSettings } from "@/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings"; import { useUpdateMCPSemanticFilterSettings } from "@/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings"; import { toast } from "@/lib/toast"; -import { Card, Col, Row, Skeleton } from "antd"; +import { Skeleton } from "@/components/ui/skeleton"; import { CircleCheck, CircleHelp, Info, Save, X } from "lucide-react"; import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { useEffect, useState } from "react"; @@ -13,6 +13,7 @@ import { FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Slider } from "@/components/ui/slider"; import { Switch } from "@/components/ui/switch"; @@ -79,6 +80,26 @@ const clampTopK = (value: number | null): number | null => const parseTopK = (raw: string, rawAsNumber: number): number | null => raw === "" || Number.isNaN(rawAsNumber) ? null : rawAsNumber; +const SaveSuccessAlert = () => { + const [dismissed, setDismissed] = useState(false); + + if (dismissed) { + return null; + } + + return ( + + + Settings saved successfully + + + + + ); +}; + export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFilterSettingsProps) { const { data, isLoading, isError, error } = useMCPSemanticFilterSettings(); const { @@ -175,7 +196,12 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi return (
    {isLoading ? ( - +
    + + + + +
    ) : isError ? ( Could not load MCP Semantic Filter settings @@ -193,17 +219,7 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi - {saveSuccess && ( - - - Settings saved successfully - - - - - )} + {saveSuccess && } {updateError && ( @@ -212,121 +228,128 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi )} - +
    {/* Left Column - Settings */} - +
    event.preventDefault()} noValidate> - - - - {({ value, onChange, onBlur, id }) => ( - commitChange(onChange, checked)} - onBlur={onBlur} - disabled={isUpdating} - /> - )} - - - - - - - - {({ value, onChange, id }) => ( - ({ - label: model.model_group, - value: model.model_group, - }))} - value={value} - onValueChange={(selected) => commitChange(onChange, selected)} - allowClear={false} - placeholder={loadingModels ? "Loading models..." : "Select embedding model"} - emptyText={loadingModels ? "Loading..." : "No embedding models available"} - disabled={isUpdating || loadingModels} - /> - )} - - - - {({ ref, value, onChange, onBlur, id }) => ( - - commitChange(onChange, parseTopK(event.target.value, event.target.valueAsNumber)) - } - onBlur={() => { - onChange(clampTopK(value)); - onBlur(); - }} - disabled={isUpdating} - /> - )} - - - - {({ value, onChange, id }) => ( -
    - + + + + {({ value, onChange, onBlur, id }) => ( + commitChange(onChange, Array.isArray(next) ? next[0] : next)} + checked={value} + onCheckedChange={(checked) => commitChange(onChange, checked)} + onBlur={onBlur} disabled={isUpdating} /> -
    - {SIMILARITY_THRESHOLD_MARKS.map((mark) => ( - - {mark.label} - - ))} + )} + + + + + + + + Configuration + + + + + {({ value, onChange, id }) => ( + ({ + label: model.model_group, + value: model.model_group, + }))} + value={value} + onValueChange={(selected) => commitChange(onChange, selected)} + allowClear={false} + placeholder={loadingModels ? "Loading models..." : "Select embedding model"} + emptyText={loadingModels ? "Loading..." : "No embedding models available"} + disabled={isUpdating || loadingModels} + /> + )} + + + + {({ ref, value, onChange, onBlur, id }) => ( + + commitChange(onChange, parseTopK(event.target.value, event.target.valueAsNumber)) + } + onBlur={() => { + onChange(clampTopK(value)); + onBlur(); + }} + disabled={isUpdating} + /> + )} + + + + {({ value, onChange, id }) => ( +
    + commitChange(onChange, Array.isArray(next) ? next[0] : next)} + disabled={isUpdating} + /> +
    + {SIMILARITY_THRESHOLD_MARKS.map((mark) => ( + + {mark.label} + + ))} +
    -
    - )} -     -     + )} + + +
    @@ -341,10 +364,10 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi
    - +
    {/* Right Column - Test Configuration */} - +
    - - +
    +
    )}
    diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/PluginSettings/PluginSettings.integration.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/PluginSettings/PluginSettings.integration.test.tsx index 998fa0ecd99..9f23fc4897c 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/PluginSettings/PluginSettings.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/PluginSettings/PluginSettings.integration.test.tsx @@ -33,7 +33,7 @@ describe("PluginSettings config payload", () => { const user = userEvent.setup(); getConfigFieldSettingMock.mockResolvedValue({ field_value: [] }); render(); - expect(await screen.findAllByText("No data")).not.toHaveLength(0); + expect(await screen.findByText("No data", { ignore: "title" })).toBeInTheDocument(); await user.click(screen.getByRole("button", { name: /add plugin/i })); fireEvent.change(await screen.findByLabelText(/Name \(identifier\)/), { target: { value: "beta" } }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RoleMappings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RoleMappings.tsx index 3750ee88183..b99cb59618f 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RoleMappings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RoleMappings.tsx @@ -1,74 +1,83 @@ import type { RoleMappings as RoleMappingsType } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; -import { Card, Divider, Table, Tag, Typography } from "antd"; +import type { ColumnDef } from "@tanstack/react-table"; +import { DataTable } from "@/components/shared/DataTable"; +import { Badge } from "@/components/ui/badge"; +import { Card, CardContent } from "@/components/ui/card"; +import { Separator } from "@/components/ui/separator"; import { Users } from "lucide-react"; import { defaultRoleDisplayNames } from "./constants"; -const { Title, Text } = Typography; + +const inlineCodeClass = "rounded-sm border border-border bg-muted px-1 py-0.5 font-mono text-xs"; + +interface RoleMappingRow { + role: string; + groups: string[]; +} export default function RoleMappings({ roleMappings }: { roleMappings: RoleMappingsType | undefined }) { if (!roleMappings) { return null; } - const roleMappingsColumns = [ + const roleMappingsColumns: ColumnDef[] = [ { - title: "Role", - dataIndex: "role", - key: "role", - render: (text: string) => {defaultRoleDisplayNames[text]}, + id: "role", + accessorKey: "role", + header: "Role", + cell: ({ row }) => {defaultRoleDisplayNames[row.original.role]}, }, { - title: "Mapped Groups", - dataIndex: "groups", - key: "groups", - render: (groups: string[]) => ( - <> - {groups.length > 0 ? ( - groups.map((group, index) => ( - + id: "groups", + accessorKey: "groups", + header: "Mapped Groups", + cell: ({ row }) => + row.original.groups.length > 0 ? ( +
    + {row.original.groups.map((group, index) => ( + {group} - - )) - ) : ( - No groups mapped - )} - - ), + + ))} +
    + ) : ( + No groups mapped + ), }, ]; return ( -
    - - Role Mappings -
    -
    -
    -
    - Group Claim -
    - {roleMappings.group_claim} -
    -
    -
    - Default Role -
    - {defaultRoleDisplayNames[roleMappings.default_role]} -
    -
    + +
    + +

    Role Mappings

    - - ({ - role, - groups, - }))} - pagination={false} - bordered - size="small" - className="w-full" - /> - +
    +
    +
    +
    Group Claim
    +
    + {roleMappings.group_claim} +
    +
    +
    +
    Default Role
    +
    + {defaultRoleDisplayNames[roleMappings.default_role]} +
    +
    +
    + + ({ + role, + groups, + }))} + getRowId={(row) => row.role} + size="compact" + /> +
    + ); } diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx index dda1b805582..504e4e45e10 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx @@ -180,19 +180,11 @@ describe("TopKeyView", () => { expect(tableViewButton).toHaveClass("bg-blue-100"); }); - it("should call setTopKeysLimit when limit is changed via Segmented control", async () => { + it("should call setTopKeysLimit when limit is changed via the segmented control", async () => { const user = userEvent.setup(); render(); - const limit10Radio = screen.getByRole("radio", { name: "10" }); - const limit10Label = limit10Radio.closest("label"); - if (limit10Label) { - await user.click(limit10Label); - } else { - // Fallback: click the div with title="10" - const limit10Div = screen.getByTitle("10"); - await user.click(limit10Div); - } + await user.click(screen.getByRole("radio", { name: "10" })); expect(mockSetTopKeysLimit).toHaveBeenCalledWith(10); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index b3c1be62e88..3f181435e09 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -3,8 +3,9 @@ import { BarChart } from "@/components/shared/charts"; import { DataTable } from "@/components/shared/DataTable"; import { IdCell, MoneyCell } from "@/components/shared/table_cells"; import { ChevronDownIcon, ChevronUpIcon } from "@heroicons/react/outline"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { SimpleTooltip } from "@/components/ui/tooltip"; -import { Segmented } from "antd"; import React, { useState } from "react"; import { formatNumberWithCommas } from "../../../../utils/dataUtils"; import { transformKeyInfo } from "../../../key_team_helpers/transform_key_info"; @@ -12,6 +13,8 @@ import { keyInfoV1Call } from "../../../networking"; import KeyInfoView from "../../../templates/key_info_view"; import { TagUsage } from "../../types"; +const TOP_KEYS_LIMITS = [5, 10, 25, 50] as const; + interface TopKeyViewProps { topKeys: any[]; teams: any[] | null; @@ -167,16 +170,22 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals return ( <>
    - setTopKeysLimit(value as number)} - /> + setTopKeysLimit(Number(limit))} + className="inline-flex w-fit items-center gap-1 rounded-lg bg-muted p-[3px]" + > + {TOP_KEYS_LIMITS.map((limit) => ( + + ))} +
    + )} +
    + + {isEditing && isGuardrailsLoading ? ( +
    Loading...
    + ) : isEditing ? ( + +
    void form.handleSubmit(onTeamUpdateSubmit)(event)}> + + + {({ ref, value, ...field }) => } + + + + {({ id, value, onChange }) => ( + + )} + + + + + {labelWithHint( + "Model Aliases", + "Map a custom alias to an underlying model. Team members can call the alias in API requests instead of the real model name.", + )} + + + + + + {({ ref, value, ...field }) => ( + + )} + + + + {({ ref, value, ...field }) => ( + + )} + + + + {({ ref, value, ...field }) => ( + + )} + + + + + Team Member Settings + + + +

    + Optional defaults applied when members join this team. All fields can be overridden per member. +

    + + + {({ id, value, onChange }) => ( + ({ + label: model, + value: model, + }))} + placeholder="Leave empty — all team models accessible to every member" + /> + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ value, onChange }) => } + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + + {({ ref, value, ...field }) => ( + + )} + + +     +     + + + {({ id, value, onChange }) => ( + onChange(next ?? null)} + /> + )} + + + + {({ ref, value, ...field }) => } + + + + {({ ref, value, ...field }) => } + + + + Metadata + + + Values are saved as text. Enter JSON for typed values, e.g. 3, true, or {'{"region": "us"}'}. + + + + + + {labelWithHint( + "Model-Specific Rate Limits", + "Set per-model TPM/RPM limits that apply across the whole team.", + )} + + {modelLimitRows.map((row, index) => ( +
    + + {({ id, value, onChange }) => ( + ({ + label: model, + value: model, + }))} + placeholder="Select model" + /> + )} + + + {({ ref, value, onChange, ...field }) => ( + ) => + onChange(event.target.value === "" ? null : Number(event.target.value)) + } + placeholder="TPM Limit" + min={0} + step={1} + /> + )} + + + {({ ref, value, onChange, ...field }) => ( + ) => + onChange(event.target.value === "" ? null : Number(event.target.value)) + } + placeholder="RPM Limit" + min={0} + step={1} + /> + )} + + +
    + ))} + +     + + + {({ ref, value, ...field }) => ( + + )} + + + + {({ ref, value, ...field }) => ( +