From e2b41286d351c6eda893c7e335d4a03896791f5e Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 00:48:27 -0400 Subject: [PATCH 01/46] fix(vertex-live): bill every modality on the /vertex_ai/live passthrough The Live passthrough builds Usage from the TEXT-modality counts alone, so audio, image and video tokens never reach the cost calculator and bill as nothing. A one-turn audio session reported 13 text and 127 audio input tokens and billed the 13; a camera session reported 1043 prompt tokens and billed 11. Reporting the full per-modality breakdown fixes it, because the shared Gemini input and output cost path already prices audio, image and video from prompt_tokens_details and completion_tokens_details. On the native-audio entry that is a 6x difference per token in both directions, which is the whole gap. Aggregation across turns is unchanged. Google charges per turn for every token in the Live session context window, current turn plus all accumulated tokens from previous turns, so the existing summing is what Vertex bills and it stays as it is. That is worth stating because the cumulative promptTokensDetails looks like a restatement of one running total, and treating it that way would under-bill a multi-turn session. See the LiveAPI context-window note on https://cloud.google.com/vertex-ai/generative-ai/pricing. Live can also name the modality carrying the rest of a turn and omit its tokenCount. Reading that absent key as zero left the tokens inside candidatesTokenCount but outside the breakdown, so real speech was charged at the text output rate. A lone unpriced entry now takes whatever the turn's declared count leaves over. Two or more cannot be told apart, so they are still left to the calculator's text remainder. Server-side toolUsePromptTokenCount is now reported in prompt_tokens_details. It is deliberately kept out of prompt_tokens: no Gemini route prices tool-use tokens, and adding them there instead suppresses the cache-overlap correction and raises the bill for no extra work. Removes _calculate_live_api_cost, whose result never reached the bill. It set kwargs["response_cost"], which the standard logging path recomputes from the ModelResponse, and on a measured audio session it returned $0.000487 against a $0.0000425 row. Now that the modality counts reach the standard calculator, keeping a second hand-rolled pricing path would only ever double-charge. The rewrite of the aggregator is arithmetically identical to what it replaced. It sums the same three counts and the same per-modality details, still takes the remaining fields from the first turn, and drops nine LIT010, one C901 and 42 basedpyright findings in the process. --- ...tex_ai_live_passthrough_logging_handler.py | 335 +++++++----------- .../test_vertex_ai_live_passthrough.py | 326 ++++++++++++----- 2 files changed, 375 insertions(+), 286 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e26f5f57532..e224f707b02 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -5,7 +5,10 @@ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough e Supports different modalities: text, audio, video, and web search. """ +from collections.abc import Mapping, Sequence from datetime import datetime +from itertools import chain +from types import MappingProxyType from typing import Any, Final from litellm._logging import verbose_proxy_logger @@ -15,8 +18,23 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( PassThroughEndpointLoggingTypedDict, ) -from litellm.types.utils import LlmProviders, ModelResponse, Usage -from litellm.utils import get_model_info +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + LlmProviders, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +_AGGREGATED_FIELDS: Final = frozenset( + { + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "promptTokensDetails", + "candidatesTokensDetails", + } +) class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @@ -48,6 +66,56 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """Return the LLM provider name.""" return LlmProviders.VERTEX_AI + @staticmethod + def _resolve_detail_counts( + details: Sequence[Mapping[str, Any]], + declared_total: object, + ) -> tuple[tuple[str, int], ...]: + """ + Pair each of one turn's ``*TokensDetails`` entries with its token count. + + Live sometimes names the modality that carries the rest of a turn without a + ``tokenCount``, and reading the absent key as zero drops those tokens from the + breakdown, so real audio ends up priced as text. A lone unpriced entry therefore takes + whatever the turn's declared count leaves over. Two or more cannot be told apart, so + they are left out and the cost calculator charges the remainder as text. + """ + priced: Final = tuple( + (str(detail.get("modality", "TEXT")), count) + for detail in details + if isinstance(count := detail.get("tokenCount"), int) + ) + unpriced: Final = tuple( + str(detail.get("modality", "TEXT")) for detail in details if not isinstance(detail.get("tokenCount"), int) + ) + if len(unpriced) != 1 or not isinstance(declared_total, int): + return priced + residual: Final = declared_total - sum(count for _, count in priced) + return priced if residual <= 0 else (*priced, (unpriced[0], residual)) + + @staticmethod + def _sum_by_modality(counts: Sequence[tuple[str, int]]) -> Mapping[str, int]: + """Total the (modality, tokenCount) pairs of one or more turns per modality.""" + return MappingProxyType({modality: sum(c for m, c in counts if m == modality) for modality, _ in counts}) + + @staticmethod + def _merged_modality_totals( + snapshots: Sequence[Mapping[str, Any]], + count_key: str, + details_key: str, + ) -> Mapping[str, int]: + """Total every turn's per-modality counts, so the breakdown adds up the way the totals do.""" + return VertexAILivePassthroughLoggingHandler._sum_by_modality( + tuple( + chain.from_iterable( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + snapshot.get(details_key) or [], snapshot.get(count_key) + ) + for snapshot in snapshots + ) + ) + ) + @staticmethod def _extract_usage_metadata_from_websocket_messages( websocket_messages: list[dict], @@ -55,175 +123,45 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Extract and aggregate usage metadata from a list of WebSocket messages. + Live emits one ``usageMetadata`` per turn and Google charges per turn for every token in + the session context window, which is the current turn's tokens plus all accumulated + tokens from previous turns, so the turns add up rather than restating each other. See + the Live API note under https://cloud.google.com/vertex-ai/generative-ai/pricing. + Args: websocket_messages: List of WebSocket messages from the Live API Returns: Dictionary containing aggregated usage metadata, or None if not found """ - all_usage_metadata: Final = [] + snapshots: Final = tuple( + message["usageMetadata"] + for message in websocket_messages + if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + ) - # Collect all usage metadata messages - for message in websocket_messages: - if isinstance(message, dict) and "usageMetadata" in message: - all_usage_metadata.append(message["usageMetadata"]) - - if not all_usage_metadata: + if not snapshots: return None - # If only one usage metadata, return it as-is - if len(all_usage_metadata) == 1: - return all_usage_metadata[0] - - # Aggregate multiple usage metadata messages - aggregated: Final[dict[str, Any]] = { - "promptTokenCount": 0, - "candidatesTokenCount": 0, - "totalTokenCount": 0, - "promptTokensDetails": [], - "candidatesTokensDetails": [], + prompt_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "promptTokenCount", "promptTokensDetails" + ) + candidate_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "candidatesTokenCount", "candidatesTokensDetails" + ) + return { + **{key: value for key, value in snapshots[0].items() if key not in _AGGREGATED_FIELDS}, + "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), + "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), + "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "promptTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 + ], + "candidatesTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in candidate_totals.items() if count > 0 + ], } - # Aggregate token counts - for usage in all_usage_metadata: - aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0) - aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0) - aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0) - - # Aggregate token details by modality - modality_totals: Final = {} - - for usage in all_usage_metadata: - # Process prompt tokens details - for detail in usage.get("promptTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["prompt"] += token_count - - # Process candidate tokens details - for detail in usage.get("candidatesTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["candidate"] += token_count - - # Convert aggregated modality totals back to details format - for modality, totals in modality_totals.items(): - if totals["prompt"] > 0: - aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]}) - if totals["candidate"] > 0: - aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]}) - - # Add any additional fields from the first usage metadata - first_usage: Final = all_usage_metadata[0] - for key, value in first_usage.items(): - if key not in aggregated: - aggregated[key] = value - - return aggregated - - @staticmethod - def _calculate_live_api_cost( - model: str, - usage_metadata: dict, - custom_llm_provider: str = "vertex_ai", - ) -> float: - """ - Calculate cost for Vertex AI Live API based on usage metadata. - - Args: - model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09") - usage_metadata: Usage metadata from the Live API response - custom_llm_provider: The LLM provider (default: "vertex_ai") - - Returns: - Total cost in USD - """ - try: - # Get model pricing information - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - - verbose_proxy_logger.debug("Vertex AI Live API model info for '%s': %s", model, model_info) - - # Check if pricing info is available - if not model_info or not model_info.get("input_cost_per_token"): - verbose_proxy_logger.error("No pricing info found for %s in local model pricing database", model) - return 0.0 - - total_cost = 0.0 - - # Extract token counts from usage metadata - prompt_token_count: Final = usage_metadata.get("promptTokenCount", 0) - candidates_token_count: Final = usage_metadata.get("candidatesTokenCount", 0) - - # Calculate base text token costs - input_cost_per_token: Final = model_info.get("input_cost_per_token", 0.0) - output_cost_per_token: Final = model_info.get("output_cost_per_token", 0.0) - - total_cost += prompt_token_count * input_cost_per_token - total_cost += candidates_token_count * output_cost_per_token - - # Handle modality-specific costs if present - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Process prompt tokens by modality - for detail in prompt_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Process candidate tokens by modality - for detail in candidates_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Handle web search costs if present - tool_use_prompt_token_count: Final = usage_metadata.get("toolUsePromptTokenCount", 0) - if tool_use_prompt_token_count > 0: - # Web search typically has a fixed cost per request - web_search_cost: Final = model_info.get("web_search_cost_per_request", 0.0) - if isinstance(web_search_cost, (int, float)) and web_search_cost > 0: - total_cost += web_search_cost - else: - # Fallback to token-based pricing for tool use - total_cost += tool_use_prompt_token_count * input_cost_per_token - - verbose_proxy_logger.debug( - f"Vertex AI Live API cost calculation - Model: {model}, " - f"Prompt tokens: {prompt_token_count}, " - f"Candidate tokens: {candidates_token_count}, " - f"Total cost: ${total_cost:.6f}" - ) - - return total_cost - - except Exception as e: - verbose_proxy_logger.error("Error calculating Vertex AI Live API cost: %s", e) - return 0.0 - @staticmethod def _create_usage_object_from_metadata( usage_metadata: dict, @@ -239,38 +177,37 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Returns: LiteLLM Usage object """ - prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) - completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) - total_tokens: Final = usage_metadata.get("totalTokenCount", 0) + prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + ) + ) + candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + ) + ) - # Create modality-specific token details if available - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Extract text tokens from details - text_prompt_tokens = 0 - text_completion_tokens = 0 - - for detail in prompt_tokens_details: - if detail.get("modality") == "TEXT": - text_prompt_tokens = detail.get("tokenCount", 0) - break - - for detail in candidates_tokens_details: - if detail.get("modality") == "TEXT": - text_completion_tokens = detail.get("tokenCount", 0) - break - - # If no text tokens found in details, use total counts - if text_prompt_tokens == 0: - text_prompt_tokens = prompt_tokens - if text_completion_tokens == 0: - text_completion_tokens = completion_tokens + prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) + completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) return Usage( - prompt_tokens=text_prompt_tokens, - completion_tokens=text_completion_tokens, - total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=prompt_by_modality.get("TEXT"), + audio_tokens=prompt_by_modality.get("AUDIO"), + image_tokens=prompt_by_modality.get("IMAGE"), + video_tokens=prompt_by_modality.get("VIDEO"), + tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=candidates_by_modality.get("TEXT"), + audio_tokens=candidates_by_modality.get("AUDIO"), + image_tokens=candidates_by_modality.get("IMAGE"), + video_tokens=candidates_by_modality.get("VIDEO"), + ), ) def vertex_ai_live_passthrough_handler( @@ -316,13 +253,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "kwargs": kwargs, } - # Calculate cost using Live API specific pricing - response_cost: Final = self._calculate_live_api_cost( - model=model, - usage_metadata=usage_metadata, - custom_llm_provider=custom_llm_provider, - ) - # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, @@ -339,8 +269,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): choices=[], ) - # Update kwargs with cost information - kwargs["response_cost"] = response_cost kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -350,10 +278,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( - f"Vertex AI Live API passthrough cost tracking - " - f"Model: {safe_model}, Cost: ${response_cost:.6f}, " - f"Prompt tokens: {usage.prompt_tokens}, " - f"Completion tokens: {usage.completion_tokens}" + "Vertex AI Live API passthrough cost tracking - Model: %s, " + "Prompt tokens: %s %s, Completion tokens: %s %s", + safe_model, + usage.prompt_tokens, + usage.prompt_tokens_details, + usage.completion_tokens, + usage.completion_tokens_details, ) return { diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index e2eb6d0b68b..3b6a548b219 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -201,88 +201,247 @@ class TestVertexAILivePassthroughLoggingHandler: assert text_prompt["tokenCount"] == 10 assert audio_prompt["tokenCount"] == 10 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_basic(self, mock_get_model_info, handler): - """Test basic cost calculation""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - } + def test_usage_carries_every_modality(self, handler): + """Regression: the Usage object reported only TEXT, so audio and image billed as nothing. + prompt_tokens must be the full count and the details must name each modality, + because the cost calculator prices audio and image from *_tokens_details. + """ usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) - - # The cost calculation may include additional factors, so we check it's reasonable - expected_min_cost = (100 * 0.000001) + (50 * 0.000002) - assert cost >= expected_min_cost - assert cost > 0 - - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_with_audio(self, mock_get_model_info, handler): - """Test cost calculation with audio tokens""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "input_cost_per_audio_token": 0.0001, - "output_cost_per_audio_token": 0.0002, - } - - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, + "promptTokenCount": 1300, + "candidatesTokenCount": 124, + "totalTokenCount": 1424, "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 80}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 13}, + {"modality": "AUDIO", "tokenCount": 127}, + {"modality": "IMAGE", "tokenCount": 1160}, ], "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 30}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 29}, + {"modality": "AUDIO", "tokenCount": 95}, ], } - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + usage = handler._create_usage_object_from_metadata( + usage_metadata=usage_metadata, model="gemini-live-2.5-flash" + ) - # Should include both text and audio costs - assert cost > 0 - assert cost > (100 * 0.000001) + ( - 50 * 0.000002 - ) # Should be higher due to audio + assert usage.prompt_tokens == 1300, "the full prompt count must survive, not just its text share" + assert usage.completion_tokens == 124 + assert usage.prompt_tokens_details.text_tokens == 13 + assert usage.prompt_tokens_details.audio_tokens == 127 + assert usage.prompt_tokens_details.image_tokens == 1160 + assert usage.completion_tokens_details.text_tokens == 29 + assert usage.completion_tokens_details.audio_tokens == 95 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" + def test_usage_sums_repeated_modality_entries(self, handler): + """A modality can appear more than once across aggregated turns; sum, don't overwrite.""" + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 0, + "promptTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 10}, + {"modality": "IMAGE", "tokenCount": 25}, + {"modality": "TEXT", "tokenCount": 5}, + ], + }, + model="gemini-live-2.5-flash", + ) + assert usage.prompt_tokens_details.image_tokens == 35 + assert usage.prompt_tokens_details.text_tokens == 5 + + NATIVE_AUDIO_MODEL = "gemini-live-2.5-flash-preview-native-audio-09-2025" + + # A four-turn native-audio session. Google charges per turn for the whole session context + # window, so the prompt side repeats the accumulated audio while the candidates side reports + # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is + # the shape Live really emits at the end of a spoken answer. + AUDIO_SESSION = ( + {"prompt": (14, 122), "candidates": (8, 20)}, + {"prompt": (21, 182), "candidates": (5, 50)}, + {"prompt": (24, 203), "candidates": (13, 27)}, + {"prompt": (24, 203), "candidates": (0, 3), "candidate_audio_token_count_missing": True}, ) - def test_calculate_cost_with_web_search(self, mock_get_model_info, handler): - """Test cost calculation with web search (tool use)""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "web_search_cost_per_request": 0.01, - } - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - "toolUsePromptTokenCount": 10, - } + @staticmethod + def _live_messages(turns): + """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" + return [{"type": "session.created", "session": {"id": "s"}}] + [ + { + "type": "response.done", + "usageMetadata": { + "promptTokenCount": sum(turn["prompt"]), + "candidatesTokenCount": sum(turn["candidates"]), + "totalTokenCount": sum(turn["prompt"]) + sum(turn["candidates"]), + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": turn["prompt"][0]}, + {"modality": "AUDIO", "tokenCount": turn["prompt"][1]}, + ], + "candidatesTokensDetails": ( + [{"modality": "AUDIO"}] + if turn.get("candidate_audio_token_count_missing") + else [ + {"modality": "TEXT", "tokenCount": turn["candidates"][0]}, + {"modality": "AUDIO", "tokenCount": turn["candidates"][1]}, + ] + ), + }, + } + for turn in turns + ] - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + @staticmethod + def _session_usage(handler, mock_logging_obj, messages, model): + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=mock_logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=model, + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + return result["result"].usage - # Should include web search cost - expected_base_cost = (100 * 0.000001) + (50 * 0.000002) - # The web search cost might be handled differently, so just check it's reasonable - assert cost >= expected_base_cost - assert cost > 0 + @classmethod + def _session_cost(cls, handler, mock_logging_obj, messages, model): + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + + usage = cls._session_usage(handler, mock_logging_obj, messages, model) + return completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + @classmethod + def _expected_session_cost(cls, turns): + from litellm.utils import get_model_info + + info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") + return ( + sum(turn["prompt"][0] for turn in turns) * info["input_cost_per_token"] + + sum(turn["prompt"][1] for turn in turns) * info["input_cost_per_audio_token"] + + sum(turn["candidates"][0] for turn in turns) * info["output_cost_per_token"] + + sum(turn["candidates"][1] for turn in turns) * info["output_cost_per_audio_token"] + ) + + def test_every_turn_of_a_session_is_billed(self, handler, mock_logging_obj): + """Google charges per turn for the whole context window, so every turn adds to the bill. + + Billing one snapshot instead gives away all the other turns: on this session the + largest single turn is well under the session total, and its share of the audio is + priced 6x the text rate, so the gap is money rather than rounding. + """ + turns = self.AUDIO_SESSION[:3] + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + widest_single_turn = max(self._expected_session_cost([turn]) for turn in turns) + assert cost > widest_single_turn, "billing one snapshot drops every other turn of the session" + + def test_audio_named_without_a_token_count_bills_at_the_audio_rate(self, handler, mock_logging_obj): + """Live can name the modality carrying the rest of a turn and omit its tokenCount. + + Reading the absent key as zero left those tokens inside candidatesTokenCount but outside + the breakdown, so the calculator charged real speech at the text output rate. At this + entry's rates the last turn's 3 audio tokens are $0.0000360 rather than $0.0000060. + """ + turns = self.AUDIO_SESSION + usage = self._session_usage(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert usage.completion_tokens_details.audio_tokens == 100, "the unpriced entry takes the turn's residual" + assert usage.completion_tokens_details.text_tokens == 26 + assert usage.completion_tokens == 126 + + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + + def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): + """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + + It is reported, not billed. Nothing in the shared Gemini input-cost path prices + tool-use tokens, and folding them into prompt_tokens here would suppress that + path's cache-overlap correction and raise the bill instead. + """ + messages = self._live_messages(self.AUDIO_SESSION[:1]) + grounded = [dict(message) for message in messages] + grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + + usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert usage.prompt_tokens_details.tool_use_tokens == 500 + + plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + + @pytest.mark.parametrize( + "label,prompt_details,candidate_details", + [ + ("text only", [("TEXT", 6)], [("TEXT", 2)]), + ("audio in", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 18)]), + ("image in", [("TEXT", 10), ("IMAGE", 258)], [("TEXT", 24)]), + ("frames in", [("TEXT", 11), ("IMAGE", 1032)], [("TEXT", 26)]), + ("audio both ways", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 29), ("AUDIO", 95)]), + ], + ) + def test_live_session_bills_each_modality_at_its_own_rate(self, handler, label, prompt_details, candidate_details): + """Every payload here is a real Vertex Live session's usageMetadata. + + Before the fix these billed the text share only, from 1x (text) to 55x under. + The expected amount is derived from the entry's own rates rather than hardcoded, + so this stays correct as prices move, and it is asserted exactly, so dropping a + modality and double-charging one both fail. + """ + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + from litellm.utils import get_model_info + + model = self.NATIVE_AUDIO_MODEL + info = get_model_info(model=model, custom_llm_provider="vertex_ai") + + text_in = info["input_cost_per_token"] + audio_in = info.get("input_cost_per_audio_token") or text_in + image_in = info.get("input_cost_per_image_token") or text_in + text_out = info["output_cost_per_token"] + audio_out = info.get("output_cost_per_audio_token") or text_out + rate_in = {"TEXT": text_in, "AUDIO": audio_in, "IMAGE": image_in} + rate_out = {"TEXT": text_out, "AUDIO": audio_out} + + expected = sum(c * rate_in[m] for m, c in prompt_details) + sum(c * rate_out[m] for m, c in candidate_details) + + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": sum(c for _, c in prompt_details), + "candidatesTokenCount": sum(c for _, c in candidate_details), + "promptTokensDetails": [{"modality": m, "tokenCount": c} for m, c in prompt_details], + "candidatesTokensDetails": [{"modality": m, "tokenCount": c} for m, c in candidate_details], + }, + model=model, + ) + + cost = completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + assert cost == pytest.approx(expected, rel=1e-9), label + + text_only = sum(c for m, c in prompt_details if m == "TEXT") * text_in + sum( + c for m, c in candidate_details if m == "TEXT" + ) * text_out + if any(m != "TEXT" for m, _ in prompt_details + candidate_details) and audio_in != text_in: + assert cost > text_only, f"{label}: non-text modalities must add cost" def test_vertex_ai_live_passthrough_handler_integration( self, handler, mock_logging_obj, sample_websocket_messages @@ -540,25 +699,24 @@ class TestVertexAILivePassthroughErrorHandling: result = handler._extract_usage_metadata_from_websocket_messages(messages) assert result is None - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_cost_calculation_with_missing_model_info(self, mock_get_model_info): - """Test cost calculation when model info is missing""" + def test_usage_without_modality_details(self): + """Older payloads carry only the totals; fall back to them rather than reporting zero.""" handler = VertexAILivePassthroughLoggingHandler() - # Mock missing model info - mock_get_model_info.return_value = {} + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + }, + model="unknown-model", + ) - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - # Should not raise an exception, should return 0 or handle gracefully - cost = handler._calculate_live_api_cost("unknown-model", usage_metadata) - assert cost == 0.0 + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 50 + assert usage.total_tokens == 150 + assert usage.prompt_tokens_details.audio_tokens is None + assert usage.prompt_tokens_details.image_tokens is None def test_handler_with_none_websocket_messages(self, mock_logging_obj): """Test handler with None websocket messages""" From da73896ec38d7cc5606d1563d87ddd7db31a4d2f Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:43:50 -0400 Subject: [PATCH 02/46] refactor(vertex-live): type the usage helpers without Any The two helpers this branch adds took Sequence[Mapping[str, Any]], which the repo forbids, and only typechecked because Any is compatible with everything. Both now take Mapping[str, object] and the raw *TokensDetails value is narrowed to its mapping entries at each of the three call sites. TypedDicts are the wrong tool here: _merged_modality_totals reads count_key and details_key as runtime strings, and the aggregation deliberately passes unknown keys straight through, so both need a mapping whose keys are not literals. The narrowing is not cosmetic. The handler's only failure path returns no result at all, so a *TokensDetails value that was not a list of objects used to raise while being read and cost the whole session its bill. --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++---- .../test_vertex_ai_live_passthrough.py | 30 +++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e224f707b02..72e8c3f7657 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -9,7 +9,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( @@ -37,6 +37,11 @@ _AGGREGATED_FIELDS: Final = frozenset( ) +def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: + """Narrow one turn's ``*TokensDetails`` value to the entries that are actually shaped like one.""" + return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -68,7 +73,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _resolve_detail_counts( - details: Sequence[Mapping[str, Any]], + details: Sequence[Mapping[str, object]], declared_total: object, ) -> tuple[tuple[str, int], ...]: """ @@ -100,7 +105,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _merged_modality_totals( - snapshots: Sequence[Mapping[str, Any]], + snapshots: Sequence[Mapping[str, object]], count_key: str, details_key: str, ) -> Mapping[str, int]: @@ -109,7 +114,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): tuple( chain.from_iterable( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - snapshot.get(details_key) or [], snapshot.get(count_key) + _detail_entries(snapshot.get(details_key)), snapshot.get(count_key) ) for snapshot in snapshots ) @@ -179,12 +184,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + _detail_entries(usage_metadata.get("promptTokensDetails")), usage_metadata.get("promptTokenCount") ) ) candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + _detail_entries(usage_metadata.get("candidatesTokensDetails")), + usage_metadata.get("candidatesTokenCount"), ) ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 3b6a548b219..ba9167e2b13 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -381,6 +381,36 @@ class TestVertexAILivePassthroughLoggingHandler: grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): + """A ``*TokensDetails`` value that is not a list of objects must not take the session down. + + The handler's only error path returns no result at all, so one odd frame used to throw + while reading it and the whole session billed nothing. The good turns still bill. + """ + turns = self.AUDIO_SESSION[:3] + messages = self._live_messages(turns) + mangled = [dict(message) for message in messages] + mangled[1]["usageMetadata"] = {**mangled[1]["usageMetadata"], "promptTokensDetails": "TEXT"} + + usage = self._session_usage(handler, mock_logging_obj, mangled, self.NATIVE_AUDIO_MODEL) + + surviving = turns[1:] + assert usage.prompt_tokens_details.audio_tokens == sum(turn["prompt"][1] for turn in surviving) + assert usage.prompt_tokens_details.text_tokens == sum(turn["prompt"][0] for turn in surviving) + assert usage.prompt_tokens == sum(sum(turn["prompt"]) for turn in turns), "the totals still cover every turn" + + direct = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 12, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 40}, "AUDIO"], + "candidatesTokensDetails": {"modality": "TEXT", "tokenCount": 12}, + }, + model=self.NATIVE_AUDIO_MODEL, + ) + assert direct.prompt_tokens_details.audio_tokens == 40, "the well-formed entry beside a bad one still counts" + assert direct.completion_tokens == 12 + @pytest.mark.parametrize( "label,prompt_details,candidate_details", [ From b051aa713a3d4c044fb9453935879f09d91eea69 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:54:13 -0400 Subject: [PATCH 03/46] fix(vertex-live): sum tool-use prompt tokens across a session's turns toolUsePromptTokenCount was the one prompt-side total not named in _AGGREGATED_FIELDS, so it rode the unknown-key pass-through and took the first frame's value while promptTokenCount, candidatesTokenCount and totalTokenCount beside it were summed. Live's frames grow over a session, so the first frame is the smallest number in the series and a grounded session under-reported its tool-use tokens by everything after turn one. It is now summed like its three neighbours. This is reporting only, and pricing these tokens is deliberately left out. Google charges tool-use prompt tokens at the input token rate, but generic_cost_per_token reads the input bill out of prompt_tokens_details and only falls back to prompt_tokens when the details carry no text or a cache hit overlaps them. Measured on the native-audio entry with 500 tool-use tokens: adding them to prompt_tokens moves an ordinary Live turn's bill by $0.0000000000, and on a turn with a cache hit it moves it by $0.0002650000 where the tokens are worth $0.0002500000, because it perturbs the cache-overlap correction. Pricing them belongs beside the modality terms in the shared input-cost path, in its own change that fixes the same latent no-op on the ordinary Gemini path. Not verified against a live capture: no Vertex Live session we have captured reported toolUsePromptTokenCount at all, so the summing convention is inferred from the three prompt-side totals that accumulate the same way. --- ...tex_ai_live_passthrough_logging_handler.py | 2 + .../test_vertex_ai_live_passthrough.py | 46 ++++++++++++++----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 72e8c3f7657..2cd336db3d6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -31,6 +31,7 @@ _AGGREGATED_FIELDS: Final = frozenset( "promptTokenCount", "candidatesTokenCount", "totalTokenCount", + "toolUsePromptTokenCount", "promptTokensDetails", "candidatesTokensDetails", } @@ -159,6 +160,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "toolUsePromptTokenCount": sum(snapshot.get("toolUsePromptTokenCount", 0) for snapshot in snapshots), "promptTokensDetails": [ {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 ], diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index ba9167e2b13..bfff52de673 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -363,22 +363,46 @@ class TestVertexAILivePassthroughLoggingHandler: cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) - def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): - """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + TOOL_USE_PER_TURN = (100, 250, 400) - It is reported, not billed. Nothing in the shared Gemini input-cost path prices - tool-use tokens, and folding them into prompt_tokens here would suppress that - path's cache-overlap correction and raise the bill instead. + def _grounded_messages(self): + """The three-turn session again, with each turn's own toolUsePromptTokenCount attached.""" + messages = self._live_messages(self.AUDIO_SESSION[:3]) + head, turns = messages[0], messages[1:] + return [head] + [ + {**message, "usageMetadata": {**message["usageMetadata"], "toolUsePromptTokenCount": tool_use}} + for message, tool_use in zip(turns, self.TOOL_USE_PER_TURN) + ] + + def test_server_side_tool_use_prompt_tokens_are_summed_over_the_session(self, handler, mock_logging_obj): + """toolUsePromptTokenCount rode the unknown-key pass-through, so it took the first turn only. + + Every other total beside it is summed across the session, and the first turn is the + smallest number in the series, so a grounded session logged far fewer tool-use tokens + than it used. This session's turns are deliberately distinct, so 750 can only come from + summing: first-turn selection gives 100, last-turn or max gives 400. """ - messages = self._live_messages(self.AUDIO_SESSION[:1]) - grounded = [dict(message) for message in messages] - grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + grounded = self._grounded_messages() usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) - assert usage.prompt_tokens_details.tool_use_tokens == 500 + assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) - plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) - grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): + """Deliberate boundary: these tokens are reported here, and priced nowhere. + + generic_cost_per_token reads the input bill out of prompt_tokens_details, and falls + back to prompt_tokens only when the details carry no text or a cache hit overlaps them, + so adding tool-use tokens to prompt_tokens is worth nothing on an ordinary Live turn and + over-charges against the cache-overlap correction when it is not. Pricing them belongs + in the shared input-cost path, beside the modality terms that already read the details. + """ + turns = self.AUDIO_SESSION[:3] + plain_cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost( + handler, mock_logging_obj, self._grounded_messages(), self.NATIVE_AUDIO_MODEL + ) + + assert plain_cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): From 22e7c7a5338a96009607c612804fe90c0dad9a1a Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 06:25:07 -0400 Subject: [PATCH 04/46] test(vertex-live): type the session helpers in the Live passthrough tests The four session helpers this PR added were unannotated. Typing them needs a name for the (text, audio) pair each turn carries, so _LiveTurn is a TypedDict rather than a Mapping union that would leave sum() over a prompt pair ill-typed, and AUDIO_SESSION is declared with it. The message list reuses list[dict[str, object]], the annotation the passthrough already uses where it collects those messages --- .../test_vertex_ai_live_passthrough.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index bfff52de673..b4d0a6c06e5 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -6,12 +6,14 @@ including the logging handler, cost tracking, and WebSocket message processing. """ import json +from collections.abc import Sequence from datetime import datetime from unittest.mock import AsyncMock, Mock, patch, MagicMock from typing import Dict, List, Any, Optional import pytest import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -22,10 +24,16 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders +from litellm.types.utils import LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth +class _LiveTurn(TypedDict): + prompt: ReadOnly[tuple[int, int]] + candidates: ReadOnly[tuple[int, int]] + candidate_audio_token_count_missing: NotRequired[ReadOnly[bool]] + + class TestVertexAILivePassthroughLoggingHandler: """Test the Vertex AI Live Passthrough Logging Handler""" @@ -257,7 +265,7 @@ class TestVertexAILivePassthroughLoggingHandler: # window, so the prompt side repeats the accumulated audio while the candidates side reports # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is # the shape Live really emits at the end of a spoken answer. - AUDIO_SESSION = ( + AUDIO_SESSION: tuple[_LiveTurn, ...] = ( {"prompt": (14, 122), "candidates": (8, 20)}, {"prompt": (21, 182), "candidates": (5, 50)}, {"prompt": (24, 203), "candidates": (13, 27)}, @@ -265,7 +273,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @staticmethod - def _live_messages(turns): + def _live_messages(turns: Sequence[_LiveTurn]) -> list[dict[str, object]]: """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" return [{"type": "session.created", "session": {"id": "s"}}] + [ { @@ -292,7 +300,12 @@ class TestVertexAILivePassthroughLoggingHandler: ] @staticmethod - def _session_usage(handler, mock_logging_obj, messages, model): + def _session_usage( + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> Usage: result = handler.vertex_ai_live_passthrough_handler( websocket_messages=messages, logging_obj=mock_logging_obj, @@ -306,7 +319,13 @@ class TestVertexAILivePassthroughLoggingHandler: return result["result"].usage @classmethod - def _session_cost(cls, handler, mock_logging_obj, messages, model): + def _session_cost( + cls, + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> float: from litellm.cost_calculator import completion_cost from litellm.types.utils import ModelResponse @@ -321,7 +340,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @classmethod - def _expected_session_cost(cls, turns): + def _expected_session_cost(cls, turns: Sequence[_LiveTurn]) -> float: from litellm.utils import get_model_info info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") From b1262b05f470f550ea0a2d879449f4fd966c4a1d Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:05:02 -0400 Subject: [PATCH 05/46] fix(gemini-live): count grounding requests so Live sessions carry their query fee Live reports grounding in the server frames and never in usageMetadata, so nothing set the counter the cost path reads and the per-query charge was missing from every grounded session. Google bills a grounded Live prompt on top of its tokens, and that fee dwarfs the token cost, so a non-zero spend check could never catch it. Both Live surfaces now read serverContent.groundingMetadata where they build usage, and reuse the chat path's own classifier so web search and Maps keep their separate SKUs rather than being counted together. Separately, a client sending turn_detection: null reached a membership test against None and took the session down with no traceback, while the branch immediately above already guards for it. Live emits grounding and usage on the same frame, verified against Vertex directly, so the realtime counter is set where usage is built. (cherry picked from commit c997436be34beb2e84f8468286b8016c195eca92) (cherry picked from commit 26c8d4822fc1c5c44fe8f72f4f57473a2ce1acbf) --- .../llms/gemini/realtime/transformation.py | 20 +++- ...tex_ai_live_passthrough_logging_handler.py | 32 ++++++- .../test_vertex_ai_live_passthrough.py | 68 +++++++++++++ .../test_gemini_realtime_transformation.py | 95 ++++++++++++++++++- 4 files changed, 212 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index c92af7de145..79985569c5f 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup return envelope.get("setup", empty_setup) +def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + """Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage. + + Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both + on the same frame, so the per-query charge is countable at the point usage is built. + """ + server_content: Final = frame.get("serverContent") + if not isinstance(server_content, Mapping): + return () + metadata: Final = server_content.get("groundingMetadata") + return (metadata,) if isinstance(metadata, Mapping) else () + + # Google bills Live transcription at an estimated 25 audio tokens/sec of input and # 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing). GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25 @@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} - elif key == "turn_detection": + elif key == "turn_detection" and value is not None: value_typed = cast(OpenAIRealtimeTurnDetection, value) if ( isinstance(value_typed, dict) @@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): {**cast(dict, message), "usageMetadata": resolved_usage_metadata}, ), ) + grounding_metadata: Final = _grounding_metadata_from_frame(message) + if grounding_metadata: + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + _chat_completion_usage, grounding_metadata + ) else: _chat_completion_usage = get_empty_usage() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 2cd336db3d6..0e2eb60704d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -43,6 +43,23 @@ def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () +def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + """Collect every ``serverContent.groundingMetadata`` a session emitted. + + Live reports grounding in the server frames, never in ``usageMetadata``, so the per-query + charge has to be counted here rather than derived from the token totals. + """ + return tuple( + metadata + for message in websocket_messages + if isinstance(message, Mapping) + for server_content in (message.get("serverContent"),) + if isinstance(server_content, Mapping) + for metadata in (server_content.get("groundingMetadata"),) + if isinstance(metadata, Mapping) + ) + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -173,6 +190,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): def _create_usage_object_from_metadata( usage_metadata: dict, model: str, + grounding_metadata: Sequence[Mapping[str, object]] = (), ) -> Usage: """ Create a LiteLLM Usage object from Live API usage metadata. @@ -180,6 +198,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Args: usage_metadata: Usage metadata from the Live API response model: The model name + grounding_metadata: Every ``serverContent.groundingMetadata`` the session emitted, so + Search and Maps grounding carry their per-query charge Returns: LiteLLM Usage object @@ -199,7 +219,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) - return Usage( + usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), @@ -217,6 +237,15 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): video_tokens=candidates_by_modality.get("VIDEO"), ), ) + if grounding_metadata: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + usage, grounding_metadata + ) + return usage def vertex_ai_live_passthrough_handler( self, @@ -264,6 +293,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), model=model, ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index b4d0a6c06e5..1815ff134aa 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -406,6 +406,74 @@ class TestVertexAILivePassthroughLoggingHandler: usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) + @staticmethod + def _grounding_frame(metadata: dict[str, object]) -> dict[str, object]: + """One server frame carrying grounding metadata, the way Live reports it.""" + return {"type": "response.done", "serverContent": {"groundingMetadata": metadata}} + + def test_web_grounding_is_counted_so_it_can_be_billed(self, handler, mock_logging_obj): + """Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames, so web_search_requests stayed unset and the cost path's only + trigger for the per-query grounding charge never fired. Google bills a grounded Live + prompt on top of its tokens, so the whole fee was missing from the bill. + """ + messages = [ + self._grounding_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_maps_grounding_is_counted_under_its_own_sku(self, handler, mock_logging_obj): + """Maps grounding is a separate SKU from web search, so it needs its own counter. + + A maps-only turn carries grounding chunks but no webSearchQueries, so counting queries + alone would report nothing and bill nothing. + """ + messages = [ + self._grounding_frame({"groundingChunks": [{"maps": {"placeId": "abc123"}}]}), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.google_maps_grounding_requests == 1 + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + + def test_an_ungrounded_session_reports_no_grounding(self, handler, mock_logging_obj): + """The counters must stay absent when no tool ran, or every session pays a grounding fee.""" + usage = self._session_usage( + handler, mock_logging_obj, self._live_messages(self.AUDIO_SESSION[:1]), self.NATIVE_AUDIO_MODEL + ) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_grounding_adds_its_query_fee_to_the_session_bill(self, handler, mock_logging_obj): + """The counter only matters if it reaches the bill, so assert against the cost, not the field. + + Same tokens either way: the difference between the two sessions is the grounding fee alone. + """ + turns = self.AUDIO_SESSION[:1] + plain = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded = self._session_cost( + handler, + mock_logging_obj, + [self._grounding_frame({"webSearchQueries": ["q"]}), *self._live_messages(turns)], + self.NATIVE_AUDIO_MODEL, + ) + + assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 2b3b6343fad..ee984cc7e1f 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,11 +1,16 @@ import json -from unittest.mock import MagicMock +from collections.abc import Mapping +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.types.llms.gemini import BidiGenerateContentServerMessage +from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents +from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2178,3 +2183,91 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra } assert usage == expected assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None + + +def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]: + """One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them.""" + from typing import Final + + server_content: Final = { + "turnComplete": True, + **({} if grounding_metadata is None else {"groundingMetadata": grounding_metadata}), + } + return { + "serverContent": server_content, + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 157, + "totalTokenCount": 176, + "promptTokensDetails": ({"modality": "TEXT", "tokenCount": 19},), + "candidatesTokensDetails": ({"modality": "AUDIO", "tokenCount": 157},), + }, + } + + +def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: + """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. + + The Usage object is local to the method, so the bridge call is the only place it is observable. + """ + from typing import Final + + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes + original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage + + def _spy(usage: Usage) -> object: + captured.append(usage) + return original(usage) + + config: Final = GeminiRealtimeConfig() + with patch.object( + LiteLLMCompletionResponsesConfig, + "_transform_chat_completion_usage_to_responses_usage", + staticmethod(_spy), + ): + config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + assert captured, "response.done must build a Usage object" + return captured[0] + + +def test_gemini_realtime_response_done_counts_web_grounding(): + """Regression: Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames on the realtime path, so web_search_requests stayed unset and the + cost path's only trigger for Google's per-query grounding charge never fired. + + Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is + built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of + detail fields and drops the rest, so the counter does not reach response.done yet. Widening + that bridge is a separate change; do not read this test as proving end-to-end billing. + """ + usage = _usage_built_for_response_done( + _grounded_live_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ) + ) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + + +def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): + """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" + usage = _usage_built_for_response_done(_grounded_live_frame(None)) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None From 9580b89bb16d9946152aa23a7a4d7d139700d403 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 23:26:31 -0400 Subject: [PATCH 06/46] fix(vertex-live): resolve a Live setup model before logging reads it A client that named a bare gateway alias logged the session as "unknown" and billed nothing, because the model was read off the raw setup frame and the extractor only yields a name when the string already contains "/models/". The rewriter qualifies that same model a few lines later for the upstream, so the supported client form, an alias, was the one that went unbilled. Resolving through the rewriter first means the real model reaches the logging object, and from there the cost map. A route with no rewriter, which is every non-Live passthrough, hands the frame over untouched. (cherry picked from commit 573982803df612fd94144e2e06dd647f8530d4e8) --- .../pass_through_endpoints.py | 23 ++++- .../test_pass_through_endpoints.py | 94 +++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ea4ede7e513..5a2f9c391a2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload +def _resolved_vertex_live_setup( + setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None +) -> Mapping[str, object]: + """ + Give the model extractor the same fully qualified path the upstream will receive. + + Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before + it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw + frame logs the session as ``unknown`` at no cost, which is precisely the supported client form + """ + setup_model: Final = setup_data.get("model") + if setup_model_rewriter is None or not isinstance(setup_model, str): + return setup_data + return {**setup_data, "model": setup_model_rewriter(setup_model)} + + def _truncated_close_reason(reason: str) -> str: """ Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character @@ -2314,7 +2330,12 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - extracted_model = _extract_model_from_vertex_ai_setup(setup_data) + # Resolve the alias first: a client may name a bare gateway model, + # which carries no "/models/" for the extractor to read, so reading + # the raw frame leaves the session logged as "unknown" and unbilled. + extracted_model = _extract_model_from_vertex_ai_setup( + _resolved_vertex_live_setup(setup_data, setup_model_rewriter) + ) if extracted_model: kwargs["model"] = extracted_model kwargs["custom_llm_provider"] = "vertex_ai-language-models" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d57bed430c1..11066d4ed38 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5030,6 +5030,100 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model(): assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" +@pytest.mark.parametrize( + "setup_model", + ["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"], +) +def test_vertex_live_setup_model_resolves_before_extraction(setup_model): + """A bare gateway alias left the session logged as ``unknown`` at zero cost. + + The model was read off the raw client frame, and the extractor only yields a name when the string + already contains ``/models/``. The rewriter qualifies it a few lines later for the upstream, so a + client that addressed the gateway the documented way, by alias, logged no model and therefore + resolved no cost-map entry. Resolving first is what puts the real name on the logging object. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + rewriter = _build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ) + setup_data = {"model": setup_model} + + resolved = _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, rewriter)) + + assert resolved == "gemini-live-2.5-flash", "an unresolved setup model logs the session as 'unknown'" + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): + """End to end through the relay: a bare alias must reach the logging object as a real model name. + + This is the call-site half of the fix. The helper tests above pass even if extraction moves back + before the rewrite, so this one drives the real websocket relay and asserts on what got logged, + which is the name the cost map is looked up by. An unbilled session logs ``unknown``. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + + upstream_ws = RecordingUpstreamWebSocket() + setup_frame = json.dumps({"setup": {"model": "gemini-live-2.5-flash"}}) + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": setup_frame}, + {"type": "websocket.disconnect"}, + ] + ) + ) + built = [] + real_logging = litellm.litellm_core_utils.litellm_logging.Logging + + def _capture(*args, **kwargs): + obj = real_logging(*args, **kwargs) + built.append(obj) + return obj + + with _patched_websocket_passthrough_environment(upstream_ws): + with patch("litellm.litellm_core_utils.litellm_logging.Logging", side_effect=_capture): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + setup_model_rewriter=_build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ), + ) + + assert built, "the relay should have built a logging object" + assert built[0].model == "gemini-live-2.5-flash", "a bare alias must not log as 'unknown'" + + +def test_vertex_live_setup_resolution_is_inert_without_a_rewriter(): + """Non-Live passthrough routes pass no rewriter, so the frame must be handed over untouched.""" + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + setup_data = {"model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"} + + assert _resolved_vertex_live_setup(setup_data, None) is setup_data + assert _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, None)) == ( + "gemini-live-2.5-flash" + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"]) async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close): From 6d71e3385b8c2ed0db61cfc1991de5614de15d83 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:12:46 -0400 Subject: [PATCH 07/46] fix(cost): carry grounding counters through the Responses usage bridge A realtime session's usage is rebuilt from its own response.done event, so a counter that does not survive the round trip is invisible to the cost path. Both directions copied a fixed allow-list, which meant a grounded Gemini Live session reported its query on the Usage object and then lost it before anything could bill it. Gemini reads the grounding counters off the input token details while Anthropic reads its own server_tool_use field, so carrying these two cannot move an Anthropic bill. Absent counters stay absent, so no provider starts paying a fee it did not incur. (cherry picked from commit ffd6c723e2a65dfff1860a913c34e41bc72ff11f) (cherry picked from commit 1590822f6893c57086b72965963d23f4f367b424) --- .../litellm_completion_transformation/transformation.py | 8 ++++++++ litellm/responses/utils.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..2b011abf0e2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2667,6 +2667,14 @@ class LiteLLMCompletionResponsesConfig: if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: input_details_dict["audio_tokens"] = prompt_details.audio_tokens + # The cost path reads the grounding counters off the input details, and a realtime + # session's usage is rebuilt from its own response.done, so dropping them here bills + # no per-query grounding fee at all. + for counter in ("web_search_requests", "google_maps_grounding_requests"): + counter_value = getattr(prompt_details, counter, None) + if counter_value is not None: + input_details_dict[counter] = counter_value + cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr( prompt_details, "cache_creation_tokens", None ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 540d492beec..7e0acd14b9c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,6 +1137,12 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), + web_search_requests=getattr( + response_api_usage.input_tokens_details, "web_search_requests", None + ), + google_maps_grounding_requests=getattr( + response_api_usage.input_tokens_details, "google_maps_grounding_requests", None + ), ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None output_tokens_details: Final[OutputTokensDetails | None] = getattr( From 6c83484065ac1b6fbed79bb9d0acea97f50a0915 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Tue, 8 Sep 2026 00:59:28 -0400 Subject: [PATCH 08/46] chore(vertex-live): drop a comment that restated the helper's docstring The call site repeated _resolved_vertex_live_setup's own docstring almost verbatim, which is the duplication the repo's comment rule exists to prevent. --- litellm/proxy/pass_through_endpoints/pass_through_endpoints.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a2f9c391a2..b66c295d1aa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2330,9 +2330,6 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - # Resolve the alias first: a client may name a bare gateway model, - # which carries no "/models/" for the extractor to read, so reading - # the raw frame leaves the session logged as "unknown" and unbilled. extracted_model = _extract_model_from_vertex_ai_setup( _resolved_vertex_live_setup(setup_data, setup_model_rewriter) ) From 4e51ff8a6dbcc6a544891f8e64b887c8cd6a8ae2 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:00:46 -0400 Subject: [PATCH 09/46] style(cost): collapse a getattr call that fits the line limit --- litellm/responses/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7e0acd14b9c..b2844834860 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,9 +1137,7 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), - web_search_requests=getattr( - response_api_usage.input_tokens_details, "web_search_requests", None - ), + web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( response_api_usage.input_tokens_details, "google_maps_grounding_requests", None ), From 228d87db6341a87bd9ae2bd1ebd76ae5e08e1a08 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:19:12 -0400 Subject: [PATCH 10/46] test(gemini-live): assert grounding counters on the emitted response.done event Replaces a patch.object spy on a static method with a read of the public return value, which also covers the usage bridge the spy ran ahead of. --- .../test_gemini_realtime_transformation.py | 66 +++++++------------ 1 file changed, 22 insertions(+), 44 deletions(-) diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index ee984cc7e1f..3eb4a70ee15 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,7 +1,7 @@ import json from collections.abc import Mapping from typing import cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -9,8 +9,6 @@ import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.types.llms.gemini import BidiGenerateContentServerMessage -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents -from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2205,40 +2203,22 @@ def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Map } -def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: - """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. - - The Usage object is local to the method, so the bridge call is the only place it is observable. - """ +def _response_done_input_details(message: Mapping[str, object]) -> Mapping[str, object]: + """The ``input_tokens_details`` a ``response.done`` event carries, read off the emitted event.""" from typing import Final - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes - original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage - - def _spy(usage: Usage) -> object: - captured.append(usage) - return original(usage) - config: Final = GeminiRealtimeConfig() - with patch.object( - LiteLLMCompletionResponsesConfig, - "_transform_chat_completion_usage_to_responses_usage", - staticmethod(_spy), - ): - config.transform_response_done_event( - message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict - BidiGenerateContentServerMessage, message - ), - current_response_id="resp_grounding", - current_conversation_id="conv_grounding", - output_items=None, - ) - assert captured, "response.done must build a Usage object" - return captured[0] + event: Final = config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + usage: Final = event["response"]["usage"] + assert usage, "response.done must carry a usage object" + return usage.get("input_tokens_details") or {} def test_gemini_realtime_response_done_counts_web_grounding(): @@ -2247,12 +2227,10 @@ def test_gemini_realtime_response_done_counts_web_grounding(): Nothing read those frames on the realtime path, so web_search_requests stayed unset and the cost path's only trigger for Google's per-query grounding charge never fired. - Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is - built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of - detail fields and drops the rest, so the counter does not reach response.done yet. Widening - that bridge is a separate change; do not read this test as proving end-to-end billing. + The counter is read off the emitted event, which is what the cost path is handed, so this covers + the grounding read and the usage bridge that carries it together """ - usage = _usage_built_for_response_done( + input_details = _response_done_input_details( _grounded_live_frame( { "webSearchQueries": ["who won the 2026 world cup final"], @@ -2261,13 +2239,13 @@ def test_gemini_realtime_response_done_counts_web_grounding(): ) ) - assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" - assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + assert input_details.get("web_search_requests") == 1, "a grounded turn must report its query" + assert input_details.get("text_tokens") == 19, "the modality breakdown must survive alongside it" def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" - usage = _usage_built_for_response_done(_grounded_live_frame(None)) + input_details = _response_done_input_details(_grounded_live_frame(None)) - assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None - assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + assert input_details.get("web_search_requests") is None + assert input_details.get("google_maps_grounding_requests") is None From e205be80df1e23a22b4a5eee5f37bc9d52cb1ce1 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 07:41:13 +0000 Subject: [PATCH 11/46] fix(proxy): enforce custom_key_update policy on /key/regenerate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 47 ++++-- .../test_key_management_endpoints.py | 150 ++++++++++++++++++ 2 files changed, 184 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 749a940de0e..106e9b8027b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -294,6 +294,35 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +async def _enforce_custom_key_update_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + data: UpdateKeyRequest, +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_update must be a coroutine") + result: Final = await hook(data) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: + changed_fields: Final = MappingProxyType( + { + field: value + for field, value in data.model_dump(exclude_unset=True).items() + if field in UpdateKeyRequest.model_fields and field != "key" + } + ) + if not changed_fields: + return None + return UpdateKeyRequest(key=key, **changed_fields) + + class _LegacyDumpable(Protocol): def dict(self) -> Mapping[str, object]: ... @@ -3066,19 +3095,7 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( - proxy_server - ) - if custom_key_update_hook is not None: - if inspect.iscoroutinefunction(custom_key_update_hook): - result: Final = await custom_key_update_hook(data) - else: - raise ValueError("user_custom_key_update must be a coroutine") - decision: Final = result.get("decision", True) - message: Final = result.get("message", "Authentication Failed - Custom Auth Rule") - if not decision: - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data) # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) @@ -5069,6 +5086,7 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import hash_token # Mirror the /key/update ownership rebind guard. See helper docstring. @@ -5116,6 +5134,9 @@ async def _execute_virtual_key_regeneration( non_default_values = {} if data is not None: + update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data) + if update_request is not None: + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request) # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) 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 65cc23ea67f..fbe790dbed5 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 @@ -24,18 +24,21 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, ProxyException, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _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.litellm_core_utils.duration_parser import duration_in_seconds 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, + _execute_virtual_key_regeneration, _get_and_validate_existing_key, _list_key_helper, _persist_deleted_verification_tokens, @@ -12011,6 +12014,153 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_hook_denies(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="3000d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ) as insert_deprecated_key, + patch( # test-quality-ok: cache eviction is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "duration must be <= 7d" + insert_deprecated_key.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert len(received_data) == 1 + assert received_data[0].key == "abc123" + assert received_data[0].duration == "3000d" + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_hook_approves(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="5d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): + mock_prisma_client = _make_regenerate_mock_prisma() + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + raise AssertionError(f"custom key update hook called with {data}") + + with ( + patch( # test-quality-ok: deterministic token setup for unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From 560e5da320bc00636057b3de3edd05d6a2da1745 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 08:02:18 +0000 Subject: [PATCH 12/46] fix(proxy): archive the old key only after regenerate validation passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 18 +++++++---------- .../test_key_management_endpoints.py | 20 +++++++++++++++---- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 106e9b8027b..97e2a4f3be5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5155,6 +5155,13 @@ async def _execute_virtual_key_regeneration( prisma_client=prisma_client, ) + await _persist_deleted_verification_tokens( + keys=[key_in_db], + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( prisma_client=prisma_client, @@ -5453,17 +5460,6 @@ async def regenerate_key_fn( if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): litellm_changed_by = None - # Save the old key record to deleted table before regeneration. - # This preserves key_alias and team_id metadata for historical spend records. - # If this fails, abort the regeneration to avoid permanently losing the - # old hash→metadata mapping. - await _persist_deleted_verification_tokens( - keys=[_key_in_db], - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - return await _execute_virtual_key_regeneration( prisma_client=prisma_client, key_in_db=_key_in_db, 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 fbe790dbed5..6ec04815420 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 @@ -11935,6 +11935,10 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: archival path is outside upperbound rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -11955,6 +11959,7 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk assert exc_info.value.status_code == 400 assert "duration" in str(exc_info.value.detail) # Rejected regenerate must not reach the DB update. + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 @@ -12037,6 +12042,10 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ) as insert_deprecated_key, + patch( # test-quality-ok: archival path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy rejection "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12063,6 +12072,7 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h assert exc_info.value.status_code == 403 assert exc_info.value.detail == "duration must be <= 7d" insert_deprecated_key.assert_not_awaited() + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 assert len(received_data) == 1 assert received_data[0].key == "abc123" @@ -12092,6 +12102,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: verify archival follows policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy approval "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12115,6 +12129,8 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho ) assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist_deleted_verification_tokens.assert_awaited_once() + assert persist_deleted_verification_tokens.call_args.kwargs["keys"] == [existing_key] assert len(received_data) == 1 @@ -13922,10 +13938,6 @@ async def test_regenerate_applies_normalized_mcp_object_permission(): "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_vector_stores_against_team", new_callable=AsyncMock, ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", - new_callable=AsyncMock, - ), patch( "litellm.proxy.management_endpoints.key_management_endpoints._execute_virtual_key_regeneration", execute_mock, From 577a4e94aa383bc48271f522a2f0e00804f749a1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:09:24 -0700 Subject: [PATCH 13/46] feat(proxy): unified custom_key_policy hook for key generate, update and regenerate Adds general_settings.custom_key_policy, one coroutine that receives the operation ("generate", "update", "regenerate"), the existing key row, the effective row as it will be written, and the raw request, and can deny with a 403. It runs after the request has been normalized and before the first DB write on /key/generate, /key/service-account/generate, /key/update, /key/bulk_update, /team/key/bulk_update and /key/{key}/regenerate. The two legacy hooks keep running unchanged on the raw request. --- .../key_management_endpoints.py | 160 +++++ litellm/proxy/proxy_server.py | 9 + .../key_management_endpoints.py | 17 + .../test_key_management_endpoints.py | 621 ++++++++++++++++++ .../test_db_overlay_remote_module_scrub.py | 1 + 5 files changed, 808 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ba43943d6de..a1ce1700f68 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -148,6 +148,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyRequest, BulkUpdateKeyResponse, BulkUpdateTeamKeysRequest, + CustomKeyPolicyRequest, FailedKeyUpdate, KeySearchWhere, SuccessfulKeyUpdate, @@ -280,6 +281,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions: class _CustomKeyHooksModule(Protocol): user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None def _custom_key_generate_hook( @@ -294,6 +296,12 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +def _custom_key_policy_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_policy + + async def _enforce_custom_key_update_policy( hook: Callable[..., Awaitable[Mapping[str, object]]] | None, data: UpdateKeyRequest, @@ -310,6 +318,113 @@ async def _enforce_custom_key_update_policy( ) +async def _enforce_custom_key_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + build_policy_request: Callable[[], CustomKeyPolicyRequest], +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_policy must be a coroutine") + result: Final = await hook(build_policy_request()) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +_KEY_UPDATE_JSON_STRING_COLUMNS: Final = frozenset({"router_settings", "budget_limits"}) + +_KEY_METADATA_REQUEST_FIELDS: Final = frozenset( + (*LiteLLM_ManagementEndpoint_MetadataFields_Premium, *LiteLLM_ManagementEndpoint_MetadataFields) +) + + +def _decode_json_string_column(column: str, value: object) -> object: + if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str): + return json.loads(value) + return value + + +def _verification_token_from_row(row: Mapping[str, object]) -> LiteLLM_VerificationToken: + org_id: Final = row["organization_id"] if "organization_id" in row else row.get("org_id") + return LiteLLM_VerificationToken.model_validate(MappingProxyType({**row, "org_id": org_id})) + + +def _effective_key_after_update( + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], +) -> LiteLLM_VerificationToken: + overlay: Final = MappingProxyType( + {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} + ) + return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + + +def _update_policy_request( + operation: Literal["update", "regenerate"], + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], + request: UpdateKeyRequest | RegenerateKeyRequest, +) -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation=operation, + existing_key=_verification_token_from_row(existing_key_row.model_dump()), + effective_key=_effective_key_after_update( + existing_key_row=existing_key_row, non_default_values=non_default_values + ), + request=request, + ) + + +def _generate_budget_windows( + budget_limits: Sequence[BudgetLimitEntry] | None, +) -> tuple[Mapping[str, object], ...] | None: + if budget_limits is None: + return None + return tuple( + MappingProxyType( + { + **window.model_dump(), + "reset_at": get_budget_reset_time(budget_duration=window.budget_duration).isoformat(), + } + ) + for window in budget_limits + ) + + +def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> LiteLLM_VerificationToken: + requested: Final = data.model_dump(exclude_unset=True, exclude_none=True) + metadata_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field in _KEY_METADATA_REQUEST_FIELDS} + ) + column_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} + ) + request_metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + expires: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None + ) + budget_reset_at: Final = ( + get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None + ) + return _verification_token_from_row( + MappingProxyType( + { + **columns, + "metadata": encrypt_callback_vars(folded_metadata), + "expires": expires, + "budget_reset_at": budget_reset_at, + "budget_limits": _generate_budget_windows(data.budget_limits), + "object_permission": None, + } + ) + ) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { @@ -1016,6 +1131,7 @@ async def _common_key_generation_helper( litellm_changed_by: str | None, team_table: LiteLLM_TeamTableCachedObj | None, ) -> GenerateKeyResponse: + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, @@ -1164,6 +1280,16 @@ async def _common_key_generation_helper( "litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)), + request=data, + ), + ) + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable _budget_id = data.budget_id if prisma_client is not None and data.soft_budget is not None: @@ -2496,6 +2622,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None = None, ) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2606,6 +2733,16 @@ async def _process_single_key_update( # Prepare update data non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=user_custom_key_policy, + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=update_key_request, + ), + ) + # Update key in database if prisma_client is None: raise HTTPException( @@ -3112,6 +3249,16 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3280,6 +3427,7 @@ async def bulk_update_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -3327,6 +3475,7 @@ async def bulk_update_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, ) successful_updates.append( @@ -3444,6 +3593,7 @@ async def bulk_update_team_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if prisma_client is None: raise HTTPException( @@ -3574,6 +3724,7 @@ async def bulk_update_team_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, existing_key_row=existing_by_token[db_token], ) @@ -5156,6 +5307,15 @@ async def _execute_virtual_key_regeneration( if new_key_alias != key_in_db.key_alias: _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="regenerate", + existing_key_row=key_in_db, + non_default_values=non_default_values, + request=data if data is not None else RegenerateKeyRequest(), + ), + ) update_data.update(non_default_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..9d258f6fd92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -924,6 +924,7 @@ def cleanup_router_config_variables(): user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -941,6 +942,7 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + user_custom_key_policy = None TEAM_METADATA_VALIDATOR_REGISTRY.set(None) TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None @@ -2365,6 +2367,7 @@ user_custom_key_generate = None _pkce_no_redis_warning_emitted: bool = False _cp_no_redis_warning_emitted: bool = False user_custom_key_update = None +user_custom_key_policy = None user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -4250,6 +4253,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Final[dict[str, tuple[str, ...]]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", @@ -5403,6 +5407,7 @@ class ProxyConfig: user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -5940,6 +5945,10 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_key_policy: Final = general_settings.get("custom_key_policy", None) + if custom_key_policy is not None: + user_custom_key_policy = get_instance_fn(value=custom_key_policy, config_file_path=config_file_path) + custom_team_metadata_validate: Final = general_settings.get("custom_team_metadata_validate", None) TEAM_METADATA_VALIDATOR_REGISTRY.set( get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 9fb5bea81e3..35abb8fb718 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -4,6 +4,9 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest +from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains @@ -123,3 +126,17 @@ class BulkUpdateTeamKeysRequest(BaseModel): if not has_key_ids and not self.all_keys_in_team: raise ValueError("Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`.") return self + + +CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] + + +class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): + """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + + model_config = ConfigDict(protected_namespaces=(), frozen=True) + + operation: CustomKeyPolicyOperation + existing_key: LiteLLM_VerificationToken | None + effective_key: LiteLLM_VerificationToken + request: GenerateKeyRequest | UpdateKeyRequest | RegenerateKeyRequest 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 000a9dc2f1a..2b8c62523b9 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 @@ -1,3 +1,4 @@ +from contextlib import ExitStack from typing import Final import json from datetime import datetime, timedelta, timezone @@ -22,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLM_VerificationToken, + LiteLLMKeyType, LitellmUserRoles, Member, ProxyException, @@ -38,6 +40,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, + _effective_key_after_update, + _effective_key_for_generate, + _enforce_custom_key_policy, _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, @@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_key_team_change, ) from litellm.proxy.proxy_server import app +from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest client = TestClient(app) @@ -12178,6 +12184,586 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +_POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" +_POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" +_POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} + + +def _seven_day_policy(received: list[CustomKeyPolicyRequest]): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + expires = policy_request.effective_key.expires + if isinstance(expires, datetime) and expires > datetime.now(timezone.utc) + timedelta(days=7): + return {"decision": False, "message": _POLICY_DENIAL_MESSAGE} + return {"decision": True} + + return policy + + +def _assert_expires_in(effective_key: LiteLLM_VerificationToken, duration: str) -> None: + expires = effective_key.expires + assert isinstance(expires, datetime) + assert expires.tzinfo is not None + expected = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration=duration)) + assert abs((expires - expected).total_seconds()) < 60 + + +def _regenerate_policy_mocks(policy, insert_deprecated_key: AsyncMock, persist: AsyncMock) -> ExitStack: + stack = ExitStack() + stack.enter_context( + patch( # test-quality-ok: deterministic token setup for the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ) + ) + stack.enter_context( + patch( # test-quality-ok: grace-period write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + insert_deprecated_key, + ) + ) + stack.enter_context( + patch( # test-quality-ok: archival write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + persist, + ) + ) + stack.enter_context( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch( # test-quality-ok: rotation callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +async def _regenerate_under_policy(mock_prisma_client, existing_key, data): + return await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_regenerate_rejects_when_custom_key_policy_denies_the_effective_expiry(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + insert_deprecated_key = AsyncMock() + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), insert_deprecated_key, persist): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="3000d")) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + insert_deprecated_key.assert_not_awaited() + persist.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key is not None + assert received[0].existing_key.token == "abc123" + assert isinstance(received[0].request, RegenerateKeyRequest) + assert received[0].request.duration == "3000d" + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_regenerate_within_custom_key_policy_rotates_the_key(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), AsyncMock(), persist): + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="5d")) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist.assert_awaited_once() + assert persist.call_args.kwargs["keys"] == [existing_key] + assert [policy_request.operation for policy_request in received] == ["regenerate"] + _assert_expires_in(received[0].effective_key, "5d") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_regenerate_without_changes_still_runs_custom_key_policy(data): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_rotation(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {"decision": False, "message": "key rotation is frozen"} + + with _regenerate_policy_mocks(freeze_rotation, AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, data) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "key rotation is frozen" + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key == existing_key + assert received[0].effective_key == existing_key + + +def _policy_existing_team_key() -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0 + ) + + +def _setup_update_key_fn_policy_mocks(monkeypatch, existing_key: LiteLLM_VerificationToken) -> AsyncMock: + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"max_budget": 50.0, "team_id": "team-a"}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=None) + ) + return mock_prisma_client + + +def _assert_update_policy_request(policy_request: CustomKeyPolicyRequest, request: UpdateKeyRequest) -> None: + assert policy_request.operation == "update" + assert policy_request.request is request + assert policy_request.existing_key is not None + assert policy_request.existing_key.max_budget == 200.0 + assert policy_request.effective_key.team_id == "team-a" + assert policy_request.effective_key.user_id == "test-user" + assert policy_request.effective_key.max_budget == 50.0 + _assert_expires_in(policy_request.effective_key, request.duration or "") + + +@pytest.mark.asyncio +async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy), # test-quality-ok: inject policy hook + ): + await update_key_fn( + request=MagicMock(), + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_update_key_fn_rejects_when_custom_key_policy_denies(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + + with patch("litellm.proxy.proxy_server.user_custom_key_policy", policy): # test-quality-ok: inject policy hook + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + _assert_expires_in(received[0].effective_key, "3000d") + + +async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data: UpdateKeyRequest, policy): + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + return await _process_single_key_update( + update_key_request=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + llm_router=None, + existing_key_row=_policy_existing_team_key(), + user_custom_key_policy=policy, + ) + + +@pytest.mark.asyncio +async def test_process_single_key_update_runs_custom_key_policy_on_the_effective_row(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + result = await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert result["max_budget"] == 50.0 + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_process_single_key_update_rejects_when_custom_key_policy_denies(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + + +@pytest.mark.asyncio +async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + ) + + existing_keys = [ + LiteLLM_VerificationToken(token="test-key-1", user_id="user-123", max_budget=None), + LiteLLM_VerificationToken(token="test-key-2", user_id="user-123", max_budget=50.0), + ] + updated_row = MagicMock() + updated_row.model_dump.return_value = {"user_id": "user-123", "max_budget": 100.0} + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=existing_keys) + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + mock_prisma_client.get_data = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + received: list[CustomKeyPolicyRequest] = [] + + async def cap_max_budget(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + max_budget = policy_request.effective_key.max_budget + if max_budget is not None and max_budget > 100: + return {"decision": False, "message": "max_budget must be 100 or less"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", cap_max_budget) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + response = await bulk_update_keys( + data=BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem(key="test-key-1", max_budget=100.0), + BulkUpdateKeyRequestItem(key="test-key-2", max_budget=500.0), + ] + ), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert [update.key for update in response.successful_updates] == ["test-key-1"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [ + ("test-key-2", "max_budget must be 100 or less") + ] + assert mock_prisma_client.update_data.await_count == 1 + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [100.0, 500.0] + assert [ + policy_request.existing_key.max_budget if policy_request.existing_key is not None else "missing" + for policy_request in received + ] == [None, 50.0] + + +def _policy_generate_prisma() -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1")) + mock_prisma.jsonify_object = MagicMock(side_effect=lambda data: json.loads(data) if isinstance(data, str) else data) + return mock_prisma + + +def _generate_policy_mocks(mock_prisma: MagicMock, generate_key_helper: AsyncMock, policy) -> ExitStack: + stack = ExitStack() + stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)) # test-quality-ok: fake DB + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", None)) # test-quality-ok: no router in test + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) # test-quality-ok: premium fields + stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) # test-quality-ok: admin + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())) # test-quality-ok: cache + stack.enter_context( + patch( # test-quality-ok: the key write must not run on a denied generate + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + generate_key_helper, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +def _generate_request(duration: str, organization_id: str | None) -> GenerateKeyRequest: + return GenerateKeyRequest( + duration=duration, + organization_id=organization_id, + guardrails=["g1"], + tags=["t1"], + soft_budget=10.0, + max_budget=20.0, + ) + + +def _assert_generate_policy_request( + policy_request: CustomKeyPolicyRequest, duration: str, organization_id: str | None +) -> None: + assert policy_request.operation == "generate" + assert policy_request.existing_key is None + assert policy_request.effective_key.org_id == organization_id + assert policy_request.effective_key.max_budget == 20.0 + assert policy_request.effective_key.metadata["guardrails"] == ["g1"] + assert policy_request.effective_key.metadata["tags"] == ["t1"] + _assert_expires_in(policy_request.effective_key, duration) + + +@pytest.mark.asyncio +async def test_generate_key_rejects_when_custom_key_policy_denies_before_any_write(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("3000d", organization_id="org-1") + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + generate_key_helper.assert_not_awaited() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "3000d", organization_id="org-1") + assert received[0].request is data + assert data.duration == "3000d" + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.organization_id == "org-1" + + +@pytest.mark.asyncio +async def test_generate_key_within_custom_key_policy_creates_the_key(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("5d", organization_id=None) + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + mock_prisma.db.litellm_budgettable.create.assert_awaited_once() + generate_key_helper.assert_awaited_once() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "5d", organization_id=None) + assert received[0].request is data + + +@pytest.mark.asyncio +async def test_service_account_generate_rejects_when_custom_key_policy_denies(): + from litellm.proxy.management_endpoints.key_management_endpoints import generate_service_account_key_fn + + mock_prisma = _policy_generate_prisma() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock()) + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + + with ( + _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)), + patch( # test-quality-ok: team lookup is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id="team-1", duration="3000d"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + generate_key_helper.assert_not_awaited() + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["generate"] + assert received[0].existing_key is None + assert received[0].effective_key.team_id == "team-1" + assert received[0].effective_key.user_id is None + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_effective_key_after_update_decodes_json_string_columns_and_keeps_omitted_fields(): + existing_key = LiteLLM_VerificationToken(token="tok", user_id="u1", team_id="team-a") + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest( + key="tok", router_settings={"num_retries": 3}, budget_limits=[{"budget_duration": "1d", "max_budget": 2.0}] + ), + existing_key_row=existing_key, + ) + assert isinstance(non_default_values["router_settings"], str) + assert isinstance(non_default_values["budget_limits"], str) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.router_settings == {"num_retries": 3} + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 2.0 + assert effective_key.budget_limits[0]["budget_duration"] == "1d" + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.team_id == "team-a" + assert effective_key.user_id == "u1" + + +@pytest.mark.asyncio +async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration(): + existing_key = LiteLLM_VerificationToken(token="tok", expires=datetime(2027, 1, 1, tzinfo=timezone.utc)) + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest(key="tok", duration="-1"), existing_key_row=existing_key + ) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.expires is None + + +def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + data = GenerateKeyRequest( + duration="5d", + organization_id="org-1", + metadata={"a": 1}, + guardrails=["g1"], + tags=["t1"], + budget_duration="1d", + max_budget=3.0, + key_type=LiteLLMKeyType.LLM_API, + ) + + effective_key = _effective_key_for_generate(data=data, now=now) + + assert effective_key.expires == now + timedelta(days=5) + assert effective_key.org_id == "org-1" + assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} + assert effective_key.max_budget == 3.0 + assert effective_key.budget_duration == "1d" + assert effective_key.budget_reset_at is not None + assert effective_key.key_type == "llm_api" + assert effective_key.allowed_routes == ["llm_api_routes"] + assert data.metadata == {"a": 1} + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.duration == "5d" + + +def test_effective_key_for_generate_without_duration_never_expires(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.expires is None + assert effective_key.budget_reset_at is None + assert effective_key.key_type == "default" + + +def _policy_request_for_generate() -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=LiteLLM_VerificationToken(token="tok"), + request=GenerateKeyRequest(), + ) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_rejects_a_sync_hook(): + def sync_hook(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": True} + + with pytest.raises(ValueError, match="user_custom_key_policy must be a coroutine"): + await _enforce_custom_key_policy(hook=sync_hook, build_policy_request=_policy_request_for_generate) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_uses_the_default_denial_message(): + async def deny(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": False} + + with pytest.raises(HTTPException) as exc_info: + await _enforce_custom_key_policy(hook=deny, build_policy_request=_policy_request_for_generate) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ @@ -18348,3 +18934,38 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): ) assert exc.value.status_code == 403 assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + keys = [_make_team_key("tok-a"), _make_team_key("tok-b")] + mock = _setup_team_keys_mocks( + monkeypatch, find_many=keys, update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}) + ) + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_tok_b(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + if policy_request.existing_key is not None and policy_request.existing_key.token == "tok-b": + return {"decision": False, "message": "tok-b is frozen"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", freeze_tok_b) + + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", key_ids=["tok-a", "tok-b"], update_fields=KeyUpdateFields(max_budget=50.0) + ) + ) + + assert [update.key for update in response.successful_updates] == ["tok-a"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [("tok-b", "tok-b is frozen")] + mock.update_data.assert_awaited_once() + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [50.0, 50.0] + assert [policy_request.effective_key.team_id for policy_request in received] == ["team-abc", "team-abc"] diff --git a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py index 100ba653f3a..0072997a0d9 100644 --- a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py +++ b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py @@ -43,6 +43,7 @@ def test_litellm_settings_callback_list_strips_remote_urls(field): "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_sso", "custom_ui_sso_sign_in_handler", ], From 83594427fca6ca448426a7fd219ce68a832ba26c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:10:10 -0700 Subject: [PATCH 14/46] fix(vertex-live): price each grounded turn's query fee on the /vertex_ai/live passthrough --- ...tex_ai_live_passthrough_logging_handler.py | 130 +++++++++++++++--- .../test_vertex_ai_live_passthrough.py | 63 ++++++++- 2 files changed, 174 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0e2eb60704d..0ac04654e30 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -7,11 +7,12 @@ Supports different modalities: text, audio, video, and web search. from collections.abc import Mapping, Sequence from datetime import datetime -from itertools import chain +from itertools import chain, pairwise from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, ) @@ -20,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrou ) from litellm.types.utils import ( CompletionTokensDetailsWrapper, + CostBreakdown, LlmProviders, ModelResponse, PromptTokensDetailsWrapper, @@ -60,6 +62,35 @@ def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[s ) +def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], ...]: + """Split a session at every ``usageMetadata`` frame; frames after the last one never got their usage.""" + closes: Final = tuple( + index + 1 + for index, message in enumerate(websocket_messages) + if isinstance(message, Mapping) and isinstance(message.get("usageMetadata"), dict) + ) + return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes))) + + +_SummedField: TypeAlias = Literal[ + "input_cost", + "output_cost", + "tool_usage_cost", + "cache_read_cost", + "cache_creation_cost", + "reasoning_cost", + "original_cost", + "discount_amount", + "margin_fixed_amount", + "margin_total_amount", +] + + +def _summed(breakdowns: Sequence[CostBreakdown], field: _SummedField) -> float | None: + values: Final = tuple(value for breakdown in breakdowns if (value := breakdown.get(field)) is not None) + return sum(values) if values else None + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -141,7 +172,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _extract_usage_metadata_from_websocket_messages( - websocket_messages: list[dict], + websocket_messages: Sequence[object], ) -> dict | None: """ Extract and aggregate usage metadata from a list of WebSocket messages. @@ -158,9 +189,11 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Dictionary containing aggregated usage metadata, or None if not found """ snapshots: Final = tuple( - message["usageMetadata"] + metadata for message in websocket_messages - if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + if isinstance(message, Mapping) + for metadata in (message.get("usageMetadata"),) + if isinstance(metadata, dict) ) if not snapshots: @@ -247,10 +280,72 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): ) return usage + def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None: + usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + if usage_metadata is None: + return None + return self._create_usage_object_from_metadata( + usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), + model=model, + ) + + def _turn_cost( + self, + turn: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> tuple[float, CostBreakdown] | None: + usage: Final = self._session_usage(turn, model) + if usage is None: + return None + cost: Final = logging_obj._response_cost_calculator( # pyright: ignore[reportPrivateUsage] # the call's own calculator keeps custom pricing and the deployment's region in step with the spend row + result=ModelResponse(model=model, usage=usage), + litellm_model_name=model, + ) + if cost is None: + return None + breakdown: Final = logging_obj.cost_breakdown + return None if breakdown is None else (cost, breakdown) + + def _session_cost( + self, + websocket_messages: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> float | None: + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) + priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) + if not priced or len(priced) != len(turn_costs): + return None + breakdowns: Final = tuple(breakdown for _, breakdown in priced) + first: Final = breakdowns[0] + total_cost: Final = sum(cost for cost, _ in priced) + logging_obj.set_cost_breakdown( + input_cost=_summed(breakdowns, "input_cost") or 0.0, + output_cost=_summed(breakdowns, "output_cost") or 0.0, + total_cost=total_cost, + cost_for_built_in_tools_cost_usd_dollar=_summed(breakdowns, "tool_usage_cost") or 0.0, + original_cost=_summed(breakdowns, "original_cost"), + discount_percent=first.get("discount_percent"), + discount_amount=_summed(breakdowns, "discount_amount"), + margin_percent=first.get("margin_percent"), + margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), + margin_total_amount=_summed(breakdowns, "margin_total_amount"), + cache_read_cost=_summed(breakdowns, "cache_read_cost"), + cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), + reasoning_cost=_summed(breakdowns, "reasoning_cost"), + service_tier=first.get("service_tier"), + data_residency=first.get("data_residency"), + vertex_location=first.get("vertex_location"), + ) + return total_cost + def vertex_ai_live_passthrough_handler( self, - websocket_messages: list[dict], - logging_obj, + websocket_messages: Sequence[object], + logging_obj: LiteLLMLoggingObj, url_route: str, start_time: datetime, end_time: datetime, @@ -274,28 +369,25 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ try: # Extract model from request body or kwargs - model: Final = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09") + requested_model: Final = kwargs.get("model") + model: Final = ( + requested_model if isinstance(requested_model, str) else "gemini-2.0-flash-live-preview-04-09" + ) custom_llm_provider: Final = kwargs.get("custom_llm_provider", "vertex_ai") verbose_proxy_logger.debug( "Vertex AI Live API model: %s, custom_llm_provider: %s", model, custom_llm_provider ) - # Extract usage metadata from WebSocket messages - usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + usage: Final = self._session_usage(websocket_messages, model) - if not usage_metadata: + if usage is None: verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages") return { "result": None, "kwargs": kwargs, } - # Create Usage object for standard LiteLLM logging - usage: Final = self._create_usage_object_from_metadata( - usage_metadata=usage_metadata, - grounding_metadata=_grounding_metadata(websocket_messages), - model=model, - ) + response_cost: Final = self._session_cost(websocket_messages, model, logging_obj) # Create a mock ModelResponse for standard logging litellm_model_response: Final = ModelResponse( @@ -306,6 +398,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): usage=usage, choices=[], ) + if response_cost is not None: + litellm_model_response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads the cost off the response's hidden params; the constructor's hidden_params kwarg is reset by pydantic kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -314,7 +408,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): import re allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") - safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" + safe_model: Final = model if allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( "Vertex AI Live API passthrough cost tracking - Model: %s, " "Prompt tokens: %s %s, Completion tokens: %s %s", diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 1815ff134aa..dbb6c4cd702 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders, Usage +from litellm.types.utils import CostBreakdown, LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth @@ -47,6 +47,7 @@ class TestVertexAILivePassthroughLoggingHandler: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @pytest.fixture @@ -474,6 +475,64 @@ class TestVertexAILivePassthroughLoggingHandler: assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def _priced_logging_obj(self) -> LiteLLMLoggingObj: + """A real logging object, since the session's price is handed to it turn by turn.""" + logging_obj = LiteLLMLoggingObj( + model=self.NATIVE_AUDIO_MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="live-session", + function_id="live", + ) + logging_obj.update_environment_variables( + model=self.NATIVE_AUDIO_MODEL, + user="u", + optional_params={}, + litellm_params={}, + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + return logging_obj + + def _billed_session( + self, handler: VertexAILivePassthroughLoggingHandler, messages: list[dict[str, object]] + ) -> tuple[float, CostBreakdown]: + logging_obj = self._priced_logging_obj() + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=self.NATIVE_AUDIO_MODEL, + custom_llm_provider="vertex_ai", + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + assert logging_obj.cost_breakdown is not None, "the session's price must reach the logging object" + return result["result"]._hidden_params["response_cost"], logging_obj.cost_breakdown + + def test_each_grounded_turn_pays_its_own_query_fee(self, handler): + """Google charges the grounding fee per grounded prompt, not per session. + + Summing the session into one usage collapsed two grounded turns into one query, so the + second question was answered for free. The bill now grows by one fee per grounded turn. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + + plain_cost, _ = self._billed_session(handler, [head, turn, turn]) + one_cost, one_breakdown = self._billed_session(handler, [head, grounding, turn, turn]) + two_cost, two_breakdown = self._billed_session(handler, [head, grounding, turn, grounding, turn]) + + fee = one_cost - plain_cost + assert fee > 0, "a grounded turn must cost more than the same tokens ungrounded" + assert two_cost - plain_cost == pytest.approx(2 * fee), "two grounded turns must pay the fee twice" + assert two_breakdown["total_cost"] == pytest.approx(two_cost) + assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. @@ -676,6 +735,7 @@ class TestVertexAILivePassthroughIntegration: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @patch( @@ -809,6 +869,7 @@ class TestVertexAILivePassthroughErrorHandling: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock def test_invalid_websocket_messages_format(self): From 710d4ae2a3aea57f78decbe856dca64e4b7e8224 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:33:31 -0700 Subject: [PATCH 15/46] fix(vertex-live): charge the fixed cost margin once per Live session --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++++++---- .../test_vertex_ai_live_passthrough.py | 23 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0ac04654e30..57ec960b032 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -314,14 +314,24 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): model: str, logging_obj: LiteLLMLoggingObj, ) -> float | None: - """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice. + + The fixed cost margin is a flat per-request fee, so the session's single spend row carries it once + rather than once per turn. + """ turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) if not priced or len(priced) != len(turn_costs): return None breakdowns: Final = tuple(breakdown for _, breakdown in priced) first: Final = breakdowns[0] - total_cost: Final = sum(cost for cost, _ in priced) + fixed_margin: Final = first.get("margin_fixed_amount") or 0.0 + duplicated_fixed_margin: Final = fixed_margin * (len(priced) - 1) + total_cost: Final = sum(cost for cost, _ in priced) - duplicated_fixed_margin + summed_margin_total: Final = _summed(breakdowns, "margin_total_amount") + margin_total_amount: Final = ( + None if summed_margin_total is None else summed_margin_total - duplicated_fixed_margin + ) logging_obj.set_cost_breakdown( input_cost=_summed(breakdowns, "input_cost") or 0.0, output_cost=_summed(breakdowns, "output_cost") or 0.0, @@ -331,8 +341,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): discount_percent=first.get("discount_percent"), discount_amount=_summed(breakdowns, "discount_amount"), margin_percent=first.get("margin_percent"), - margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), - margin_total_amount=_summed(breakdowns, "margin_total_amount"), + margin_fixed_amount=first.get("margin_fixed_amount"), + margin_total_amount=margin_total_amount, cache_read_cost=_summed(breakdowns, "cache_read_cost"), cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), reasoning_cost=_summed(breakdowns, "reasoning_cost"), diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index dbb6c4cd702..70c2fda369b 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -13,6 +13,7 @@ from typing import Dict, List, Any, Optional import pytest import httpx +import litellm from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -533,6 +534,28 @@ class TestVertexAILivePassthroughLoggingHandler: assert two_breakdown["total_cost"] == pytest.approx(two_cost) assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_the_fixed_cost_margin_is_charged_once_per_session(self, handler): + """A fixed cost margin is a flat per-request fee, and a Live session is one spend row. + + Pricing each turn on its own applied the fixed margin per turn, so a two-turn session paid it + twice. The session now carries the fixed margin once no matter how many turns it billed. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + messages = [head, grounding, turn, grounding, turn] + + plain_cost, _ = self._billed_session(handler, messages) + + fixed_amount = 0.01 + with patch.object(litellm, "cost_margin_config", {"vertex_ai": {"fixed_amount": fixed_amount}}): + margined_cost, breakdown = self._billed_session(handler, messages) + + assert margined_cost - plain_cost == pytest.approx( + fixed_amount + ), "a two-turn session must add the fixed margin once, not once per billed turn" + assert breakdown["margin_fixed_amount"] == pytest.approx(fixed_amount) + assert breakdown["margin_total_amount"] == pytest.approx(fixed_amount) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. From fc62df33c2a0b049e8b48baad9bf97cd87aa57d3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:42:20 -0700 Subject: [PATCH 16/46] fix(proxy): resolve rotation and permission fields before the key policy and pin the effective-row contract --- .../key_management_endpoints.py | 38 +++++++---- .../key_management_endpoints.py | 13 +++- .../test_key_management_endpoints.py | 68 ++++++++++++++++++- 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a1ce1700f68..841c22da050 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -359,7 +359,9 @@ def _effective_key_after_update( overlay: Final = MappingProxyType( {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} ) - return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + return _verification_token_from_row( + MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None}) + ) def _update_policy_request( @@ -381,7 +383,7 @@ def _update_policy_request( def _generate_budget_windows( budget_limits: Sequence[BudgetLimitEntry] | None, ) -> tuple[Mapping[str, object], ...] | None: - if budget_limits is None: + if not budget_limits: return None return tuple( MappingProxyType( @@ -402,15 +404,20 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite column_fields: Final = MappingProxyType( {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} ) - request_metadata: Final = data.metadata or MappingProxyType({}) - folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict - columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place expires: Final = ( now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None ) budget_reset_at: Final = ( get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None ) + key_rotation_at: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval)) + if data.auto_rotate and data.rotation_interval + else None + ) return _verification_token_from_row( MappingProxyType( { @@ -418,6 +425,7 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite "metadata": encrypt_callback_vars(folded_metadata), "expires": expires, "budget_reset_at": budget_reset_at, + "key_rotation_at": key_rotation_at, "budget_limits": _generate_budget_windows(data.budget_limits), "object_permission": None, } @@ -3249,16 +3257,6 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) - await _enforce_custom_key_policy( - hook=_custom_key_policy_hook(proxy_server), - build_policy_request=lambda: _update_policy_request( - operation="update", - existing_key_row=existing_key_row, - non_default_values=non_default_values, - request=data, - ), - ) - # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3278,6 +3276,16 @@ async def update_key_fn( existing_key_alias=existing_key_row.key_alias, ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + if prisma_client is None: raise Exception("Not connected to DB!") diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 35abb8fb718..63bbaa5ba4e 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict @@ -128,11 +128,18 @@ class BulkUpdateTeamKeysRequest(BaseModel): return self -CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] +CustomKeyPolicyOperation: TypeAlias = Literal["generate", "update", "regenerate"] class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): - """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + """What `general_settings.custom_key_policy` receives. + + `effective_key` is the verification token row as it will be written: the existing row overlaid with the + requested changes, with `duration` resolved to `expires` and `budget_duration` to `budget_reset_at`. Values the + proxy fills in after the policy stay at their defaults: `token`, `key_name`, `created_by`, `updated_by` and the + soft-budget `budget_id` on generate, the rotated token on regenerate, and the `object_permission` relation on + every operation (`object_permission_id` is set; read `request.object_permission` for the requested change). + """ model_config = ConfigDict(protected_namespaces=(), frozen=True) 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 2b8c62523b9..b014e4428fd 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 @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ResetSpendRequest, UpdateKeyRequest, ) +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -12362,7 +12363,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) received: list[CustomKeyPolicyRequest] = [] policy = _seven_day_policy(received) - data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + data = UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0, auto_rotate=True, rotation_interval="30d" + ) with ( patch( # test-quality-ok: cache eviction is outside the policy path @@ -12381,6 +12384,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client.update_data.assert_awaited_once() assert len(received) == 1 _assert_update_policy_request(received[0], data) + key_rotation_at = received[0].effective_key.key_rotation_at + assert key_rotation_at is not None + assert abs(key_rotation_at - (datetime.now(timezone.utc) + timedelta(days=30))) < timedelta(seconds=60) @pytest.mark.asyncio @@ -12695,6 +12701,23 @@ async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration assert effective_key.expires is None +def test_effective_key_after_update_swaps_the_object_permission_id_and_drops_the_stale_relation(): + existing_key = LiteLLM_VerificationToken( + token="tok", + object_permission_id="op-old", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-old", mcp_servers=["old"]), + ) + + effective_key = _effective_key_after_update( + existing_key_row=existing_key, non_default_values={"object_permission_id": "op-new"} + ) + + assert effective_key.object_permission_id == "op-new" + assert effective_key.object_permission is None + assert existing_key.object_permission is not None + assert existing_key.object_permission.mcp_servers == ["old"] + + def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): now = datetime(2026, 1, 1, tzinfo=timezone.utc) data = GenerateKeyRequest( @@ -12705,12 +12728,21 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat tags=["t1"], budget_duration="1d", max_budget=3.0, + budget_limits=[{"budget_duration": "1d", "max_budget": 5.0}], + auto_rotate=True, + rotation_interval="30d", + object_permission={"mcp_servers": ["srv"]}, key_type=LiteLLMKeyType.LLM_API, ) effective_key = _effective_key_for_generate(data=data, now=now) assert effective_key.expires == now + timedelta(days=5) + assert effective_key.key_rotation_at == now + timedelta(days=30) + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 5.0 + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.object_permission is None assert effective_key.org_id == "org-1" assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} assert effective_key.max_budget == 3.0 @@ -12722,6 +12754,18 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat assert data.guardrails == ["g1"] assert data.tags == ["t1"] assert data.duration == "5d" + assert data.budget_limits is not None + assert data.budget_limits[0].reset_at is None + assert data.object_permission is not None + assert data.object_permission.mcp_servers == ["srv"] + + +def test_effective_key_for_generate_stores_no_budget_windows_for_an_empty_list(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(budget_limits=[]), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.budget_limits is None def test_effective_key_for_generate_without_duration_never_expires(): @@ -12731,6 +12775,7 @@ def test_effective_key_for_generate_without_duration_never_expires(): assert effective_key.expires is None assert effective_key.budget_reset_at is None + assert effective_key.key_rotation_at is None assert effective_key.key_type == "default" @@ -12764,6 +12809,27 @@ async def test_enforce_custom_key_policy_uses_the_default_denial_message(): assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_allows_when_the_decision_is_missing(): + received: list[CustomKeyPolicyRequest] = [] + + async def no_decision(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {} + + await _enforce_custom_key_policy(hook=no_decision, build_policy_request=_policy_request_for_generate) + + assert len(received) == 1 + assert received[0].operation == "generate" + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_never_builds_the_request_without_a_hook(): + await _enforce_custom_key_policy( + hook=None, build_policy_request=lambda: pytest.fail("policy request built without a hook") + ) + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From a159c7d98eb786a5bc7285dee0984512b6965a21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:38:23 -0700 Subject: [PATCH 17/46] fix(proxy): write the requested object permission row only after the key policy allows the update --- .../key_management_endpoints.py | 38 ++-- .../test_key_management_endpoints.py | 164 +++++++++++++++++- 2 files changed, 178 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 841c22da050..ad2ab58d357 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2464,12 +2464,6 @@ async def prepare_key_update_data( # sentinel for Json? columns, so store the JSON literal null non_default_values["budget_limits"] = json.dumps(None) - if "object_permission" in non_default_values: - non_default_values = await _handle_update_object_permission( - data_json=non_default_values, - existing_key_row=existing_key_row, - ) - _metadata: Final = existing_key_row.metadata or {} # validate model_max_budget @@ -2490,13 +2484,12 @@ async def prepare_key_update_data( async def _handle_update_object_permission( data_json: dict, existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, ) -> dict: - """ - Handle the update of object permission. - """ - from litellm.proxy.proxy_server import prisma_client + """Persist the requested object permission row and swap it for its id, only after the key policy allowed the write.""" + if "object_permission" not in data_json: + return data_json - # Use the common helper to handle the object permission update object_permission_id: Final = await handle_update_object_permission_common( data_json=data_json, existing_object_permission_id=existing_key_row.object_permission_id, @@ -2758,7 +2751,12 @@ async def _process_single_key_update( detail={"error": "Database not connected"}, ) - _data: Final = {**non_default_values, "token": update_key_request.key} + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) + _data: Final = {**update_values, "token": update_key_request.key} response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict "Mapping[str, object] | None", await prisma_client.update_data(token=update_key_request.key, data=_data), @@ -3289,18 +3287,23 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( await _update_key_row_with_soft_budget( prisma_client=prisma_client, key=key, data=data, - non_default_values=non_default_values, + non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, ) if "soft_budget" in data.model_fields_set - else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key})) + else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key})) ) # Delete - key from cache, since it's been updated! @@ -5324,7 +5327,12 @@ async def _execute_virtual_key_regeneration( request=data if data is not None else RegenerateKeyRequest(), ), ) - update_data.update(non_default_values) + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=key_in_db, + prisma_client=prisma_client, + ) + update_data.update(update_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash, 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 b014e4428fd..82458b5dfa3 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 @@ -1,3 +1,4 @@ +from collections.abc import Mapping from contextlib import ExitStack from typing import Final import json @@ -18,6 +19,7 @@ from litellm.proxy._types import ( GenerateKeyRequest, NewUserRequest, LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionBase, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj, @@ -1035,7 +1037,7 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_existing_permission(monkeypatch): +async def test_key_update_object_permissions_existing_permission(): """ Test updating object permissions when a key already has an existing object_permission_id. @@ -1055,9 +1057,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Mock existing key with object_permission_id existing_key_row = LiteLLM_VerificationToken( @@ -1097,6 +1097,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row, + prisma_client=mock_prisma_client, ) # Verify the object_permission was removed from data_json and object_permission_id was set @@ -1111,7 +1112,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_no_existing_permission(monkeypatch): +async def test_key_update_object_permissions_no_existing_permission(): """ Test creating object permissions when a key has no existing object_permission_id. @@ -1131,9 +1132,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_no_perm = LiteLLM_VerificationToken( token="test_token_hash_2", @@ -1164,6 +1163,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_no_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -1174,7 +1174,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) @pytest.mark.asyncio -async def test_key_update_object_permissions_missing_permission_record(monkeypatch): +async def test_key_update_object_permissions_missing_permission_record(): """ Test creating object permissions when existing object_permission_id record is not found. @@ -1194,9 +1194,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_missing_perm = LiteLLM_VerificationToken( token="test_token_hash_3", @@ -1227,6 +1225,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_missing_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -12470,6 +12469,153 @@ async def test_process_single_key_update_rejects_when_custom_key_policy_denies() assert [policy_request.operation for policy_request in received] == ["update"] +_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy" + + +def _record_object_permission_writes(mock_prisma_client: AsyncMock, events: list[str]) -> None: + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + + async def upsert(**_kwargs: object) -> MagicMock: + events.append("permission row upsert") + return MagicMock(object_permission_id=_OBJECT_PERMISSION_ID_AFTER_POLICY) + + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(side_effect=upsert) + + +def _recording_policy(events: list[str], allowed: bool): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + events.append("policy") + return {"decision": allowed, "message": "key max_budget must be 1000 or less"} + + return policy + + +def _assert_permission_row_written_after_policy(events: list[str], written: Mapping[str, object]) -> None: + assert events == ["policy", "permission row upsert"] + assert written["object_permission_id"] == _OBJECT_PERMISSION_ID_AFTER_POLICY + assert "object_permission" not in written + + +def _assert_permission_row_untouched(mock_prisma_client: AsyncMock, events: list[str]) -> None: + assert events == ["policy"] + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + +def _update_with_object_permission(max_budget: float) -> UpdateKeyRequest: + return UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, + max_budget=max_budget, + object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"]), + ) + + +def _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed: bool) -> tuple[AsyncMock, list[str]]: + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", _recording_policy(events, allowed)) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", AsyncMock() + ) + return mock_prisma_client, events + + +async def _update_key_fn_with_object_permission(max_budget: float): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + return await update_key_fn( + request=MagicMock(), + data=_update_with_object_permission(max_budget=max_budget), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +async def test_update_key_fn_writes_the_object_permission_row_only_after_the_policy_allows(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=True) + + await _update_key_fn_with_object_permission(max_budget=50.0) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_update_key_fn_denied_by_the_policy_leaves_the_object_permission_row_untouched(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=False) + + with pytest.raises(ProxyException) as exc_info: + await _update_key_fn_with_object_permission(max_budget=5000.0) + + assert str(exc_info.value.code) == "403" + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_process_single_key_update_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=50.0), _recording_policy(events, allowed=True) + ) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_process_single_key_update_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=5000.0), _recording_policy(events, allowed=False) + ) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_regenerate_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=50.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=True), AsyncMock(), AsyncMock()): + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + _assert_permission_row_written_after_policy( + events, mock_prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"] + ) + + +@pytest.mark.asyncio +async def test_regenerate_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=5000.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=False), AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + @pytest.mark.asyncio async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys From 0429c645043e363ae2a87700a858349657850f6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:22:31 -0700 Subject: [PATCH 18/46] fix(proxy): keep an empty duration out of the legacy update hook on key regenerate --- .../key_management_endpoints.py | 7 ++- .../test_key_management_endpoints.py | 57 ++++++++++++++++++- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ad2ab58d357..8aaaba09e22 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -433,12 +433,17 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite ) +_EMPTY_DURATION_MEANS_UNCHANGED: Final = frozenset({"duration", "budget_duration"}) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { field: value for field, value in data.model_dump(exclude_unset=True).items() - if field in UpdateKeyRequest.model_fields and field != "key" + if field in UpdateKeyRequest.model_fields + and field != "key" + and not (field in _EMPTY_DURATION_MEANS_UNCHANGED and value == "") } ) if not changed_fields: 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 82458b5dfa3..df58424095f 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 @@ -12142,7 +12142,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho @pytest.mark.asyncio -@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +@pytest.mark.parametrize( + "data", + [None, RegenerateKeyRequest(), RegenerateKeyRequest(duration=""), RegenerateKeyRequest(budget_duration="")], +) async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): mock_prisma_client = _make_regenerate_mock_prisma() @@ -12184,6 +12187,58 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_hides_the_untouched_modal_expiry_from_the_custom_key_update_hook(): + mock_prisma_client = _make_regenerate_mock_prisma() + untouched_modal_body = RegenerateKeyRequest( + key_alias=None, max_budget=None, tpm_limit=None, rpm_limit=None, duration="", grace_period="" + ) + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration is not None and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for the untouched modal body + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=untouched_modal_body, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + assert "duration" not in received_data[0].model_fields_set + assert received_data[0].model_fields_set >= {"key", "key_alias", "max_budget", "tpm_limit", "rpm_limit"} + + _POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" _POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" _POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} From cccff657cf376843a1e9c732e9ec33c46183e01e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:35:19 -0700 Subject: [PATCH 19/46] fix(guardrails): scan the Anthropic top-level system prompt and tool_use arguments --- .../chat/guardrail_translation/handler.py | 224 +++++++++++++--- .../test_anthropic_guardrail_handler.py | 248 ++++++++++++++++-- 2 files changed, 413 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..eb16278c9bd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -103,9 +103,24 @@ class ToolResultBlockTextTarget: block_idx: int -InputWriteBackTarget = ( - MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget -) +@dataclass(frozen=True, slots=True) +class SystemStringTarget: + pass + + +@dataclass(frozen=True, slots=True) +class SystemBlockTextTarget: + block_idx: int + + +@dataclass(frozen=True, slots=True) +class ToolUseInputTarget: + msg_idx: int + content_idx: int + + +MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget +InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: @@ -146,10 +161,17 @@ class ScannedText: target: InputWriteBackTarget +@dataclass(frozen=True, slots=True) +class ScannedToolCall: + tool_call: ChatCompletionToolCallChunk + target: ToolUseInputTarget + + @dataclass(frozen=True, slots=True) class ExtractedInput: scanned: tuple[ScannedText, ...] images: tuple[str, ...] + tool_calls: tuple[ScannedToolCall, ...] = () EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) @@ -161,6 +183,76 @@ class _ToolCallShape: arguments: str +def _is_client_tool_use(block: Mapping[str, object]) -> bool: + return ( + block.get("type") == "tool_use" + and isinstance(block.get("id"), str) + and isinstance(block.get("name"), str) + and isinstance(block.get("input"), Mapping) + ) + + +def _write_back_system_block(system: object, block_idx: int, response: str) -> None: + if not isinstance(system, list): + return + text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text") + if block_idx < len(text_blocks): + text_blocks[block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + + +def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None: + content: Final = message.get("content", None) + if content is None: + return + match target: + case MessageContentTarget(): + if isinstance(content, str): + message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place + case ContentBlockTextTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultStringTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["content"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): + if isinstance(content, list): + content[content_idx]["content"][block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case _: + assert_never(target) + + +def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: + content: Final = message.get("content", None) + block: Final = content[target.content_idx] if isinstance(content, list) else None + if not isinstance(block, dict): + return + try: + rewritten_input: Final = json.loads(shape.arguments) + except json.JSONDecodeError: + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + if not isinstance(rewritten_input, dict): + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place + if shape.name is not None and shape.name != block.get("name"): + block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place + + @dataclass(frozen=True, slots=True) class _SSEFieldRewrite: """One field of one nested section of a buffered SSE event, rewritten.""" @@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation): skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply) scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) - # Exclude only the trusted top-level prompt. In-sequence system entries are untrusted - # and must stay aligned with texts_to_check for positional masking. When the top-level - # prompt is included, the pre-existing count mismatch disables positional masking. + # The top-level prompt is translated on its own below so it can be hoisted in front of + # any mid-turn system entries and scanned first, aligned with that structured position. translation_source: Final = { # mutable-ok: API message payload key: value for key, value in data.items() if key != "system" } @@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation): ] ) - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tool calls + top_level_system_scanned: Final = ( + () + if hoisted_system_message is None or scan_only_tool_results + else self._extract_top_level_system_text(hoisted_system_message) + ) extracted: Final = tuple( self._extract_input_text_and_images( message=message, @@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation): ) for msg_idx, message in enumerate(messages) ) - scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned) + scanned: Final = ( + *top_level_system_scanned, + *(item for one_message in extracted for item in one_message.scanned), + ) texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str] images_to_check: Final = [ image for one_message in extracted for image in one_message.images ] # mutable-ok: GenericGuardrailAPIInputs takes list[str] + scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls) + tool_calls_to_check: Final = [ + item.tool_call for item in scanned_tool_calls + ] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk] + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check if tools_to_check: inputs["tools"] = tools_to_check original_structured_messages: Final = structured_messages @@ -572,10 +678,16 @@ class AnthropicMessagesHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( - messages=messages, + data=data, responses=guardrailed_texts, scanned=scanned, ) + self._apply_guardrail_tool_calls_to_input( + messages=messages, + scanned_tool_calls=scanned_tool_calls, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -598,6 +710,19 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload return hoisted[0] if hoisted else None + @staticmethod + def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]: + content: Final = hoisted_system_message.get("content") + if isinstance(content, str): + return (ScannedText(content, SystemStringTarget()),) + if not isinstance(content, list): + return () + return tuple( + ScannedText(text_str, SystemBlockTextTarget(block_idx)) + for block_idx, block in enumerate(content) + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + ) + @staticmethod def _openai_system_message_to_anthropic( message: Mapping[str, object], @@ -852,9 +977,25 @@ class AnthropicMessagesHandler(BaseTranslation): for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) ) + tool_use_blocks: Final = ( + () + if scan_only_tool_results + else tuple( + (content_idx, content_item) + for content_idx, content_item in enumerate(content) + if isinstance(content_item, dict) and _is_client_tool_use(content_item) + ) + ) return ExtractedInput( scanned=tuple(item for block in blocks for item in block.scanned), images=tuple(image for block in blocks for image in block.images), + tool_calls=tuple( + ScannedToolCall( + tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx), + target=ToolUseInputTarget(msg_idx, content_idx), + ) + for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks) + ), ) @classmethod @@ -940,43 +1081,48 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: Sequence[_WritableMessage], + data: dict, responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: """ - Apply guardrail responses back to input messages. + Apply guardrail responses back to the top-level system prompt and the input messages. """ + messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () for item, guardrail_response in zip(scanned, responses): - target = item.target - message = messages[target.msg_idx] - content = message.get("content", None) - if content is None: - continue - - match target: - case MessageContentTarget(): - if isinstance(content, str): - message["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ContentBlockTextTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["text"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultStringTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): - if isinstance(content, list): - content[content_idx]["content"][block_idx]["text"] = ( + match item.target: + case SystemStringTarget(): + if isinstance(data.get("system"), str): + data["system"] = ( guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place ) + case SystemBlockTextTarget(block_idx=block_idx): + _write_back_system_block(data.get("system"), block_idx, guardrail_response) + case ( + MessageContentTarget() + | ContentBlockTextTarget() + | ToolResultStringTarget() + | ToolResultBlockTextTarget() as message_target + ): + _write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response) case _: - assert_never(target) + assert_never(item.target) + + @staticmethod + def _apply_guardrail_tool_calls_to_input( + messages: Sequence[_WritableMessage], + scanned_tool_calls: tuple[ScannedToolCall, ...], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: object, + ) -> None: + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tuple(item.tool_call for item in scanned_tool_calls) + ) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): + if before != after: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after) async def process_output_response( self, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 9fe56f4dc65..4b47b7d8c4a 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -635,14 +635,19 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.inputs is not None - assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + assert guardrail.inputs["texts"] == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + ] structured = guardrail.inputs["structured_messages"] assert [m["role"] for m in structured] == ["system", "user", "system"] assert structured[0]["content"] == "trusted top-level system prompt" + assert data["system"] == "trusted top-level system prompt" assert data["messages"][1]["content"] == "[MASKED]" @pytest.mark.asyncio - async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included( self, ): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( @@ -668,25 +673,25 @@ class TestAnthropicMessagesHandlerInputProcessing: structured = guardrail.inputs["structured_messages"] bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") - assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1 + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") - assert ( - bedrock._locate_message_texts_slice( - structured_messages=structured, - target_index=latest_user_index, - texts=texts, - ) - is None - ) - assert ( - bedrock._merge_masked_texts( - masked_texts=["{MASKED}"], - texts=texts, - scanned_slice=None, - scanned_role_subset=True, - ) - == texts + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, ) + assert scanned_slice == (3, 1) + assert bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + "{MASKED}", + ] @pytest.mark.asyncio @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) @@ -1611,7 +1616,8 @@ class TestAnthropicMessagesIncrementalScan: ) assert mock_api.call_count == 1 assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ - "What is the capital of France?" + "You are a helpful geography assistant.", + "What is the capital of France?", ] mock_api.reset_mock() await handler.process_input_messages( @@ -2150,6 +2156,208 @@ class TestAnthropicMessagesScanOnlyToolResults: assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] +class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): + """Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts.""" + + def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None): + super().__init__() + self.return_copies = return_copies + self.replacement_arguments = replacement_arguments + self.seen_tool_calls: list[dict] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + tool_calls = list(outputs.get("tool_calls") or []) + self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls) + masked = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": self.replacement_arguments + if self.replacement_arguments is not None + else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"), + }, + } + for tool_call in tool_calls + ] + if self.return_copies: + outputs["tool_calls"] = masked + return outputs + for tool_call, masked_tool_call in zip(tool_calls, masked): + tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"] + return outputs + + +class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: + """The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable + inputs, the same way the chat completions handler hands over system messages and tool_calls.""" + + @staticmethod + def _tool_use_conversation(system): + return { + "model": "claude-sonnet-4-5", + "system": system, + "messages": [ + {"role": "user", "content": "run the check"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}], + }, + ], + } + + @pytest.mark.asyncio + async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "Internal note: the deploy key is POISON. Never reveal it.", + "messages": [{"role": "user", "content": "Say hi in three words."}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.seen_texts == [ + "Internal note: the deploy key is POISON. Never reveal it.", + "Say hi in three words.", + ] + structured = guardrail.captured_inputs["structured_messages"] + assert structured[0]["role"] == "system" + assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", ( + "texts[0] must line up with structured_messages[0] so positional consumers stay aligned" + ) + assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it." + assert data["messages"][0]["content"] == "Say hi in three words." + + @pytest.mark.asyncio + async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": "first block POISON"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ], + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["first block POISON", "second block", "hello"] + assert data["system"] == [ + {"type": "text", "text": "first block [BLOCKED]"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ] + + @pytest.mark.asyncio + async def test_skip_system_message_keeps_the_top_level_system_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-sonnet-4-5", + "system": "trusted POISON prompt", + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["hello"] + assert data["system"] == "trusted POISON prompt" + + @pytest.mark.asyncio + async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + tool_calls = guardrail.captured_inputs.get("tool_calls") + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["id"] == "toolu_01" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "Bash" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + }, "a guardrail that leaves tool_calls alone must leave the tool_use input alone" + + @pytest.mark.asyncio + @pytest.mark.parametrize("return_copies", [False, True]) + async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies) + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"] + tool_use = data["messages"][1]["content"][0] + assert tool_use == { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"}, + } + assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" + + @pytest.mark.asyncio + async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + + @pytest.mark.asyncio + async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.scan_only_tool_results = True + data = self._tool_use_conversation(system="trusted POISON prompt") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"] + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tool_calls") is None + assert data["system"] == "trusted POISON prompt" + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page" + + class TestStructuredWriteBackKeepsToolResults: """A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103).""" From f7e9277032651786d493c72225e571d28f40c136 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:13:17 -0700 Subject: [PATCH 20/46] fix(guardrails): validate rewritten tool_use arguments with a typed adapter --- .../chat/guardrail_translation/handler.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index eb16278c9bd..b438d168b52 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -20,6 +20,7 @@ from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger @@ -188,7 +189,7 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool: block.get("type") == "tool_use" and isinstance(block.get("id"), str) and isinstance(block.get("name"), str) - and isinstance(block.get("input"), Mapping) + and isinstance(block.get("input"), dict) ) @@ -229,22 +230,19 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge assert_never(target) +_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: content: Final = message.get("content", None) block: Final = content[target.content_idx] if isinstance(content, list) else None if not isinstance(block, dict): return try: - rewritten_input: Final = json.loads(shape.arguments) - except json.JSONDecodeError: + rewritten_input: Final = _TOOL_USE_INPUT_ADAPTER.validate_json(shape.arguments) + except ValidationError: verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", - block.get("id"), - ) - return - if not isinstance(rewritten_input, dict): - verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + "Anthropic Messages: guardrail returned arguments that are not a JSON object for tool_use %s; keeping its input", block.get("id"), ) return @@ -1113,11 +1111,11 @@ class AnthropicMessagesHandler(BaseTranslation): messages: Sequence[_WritableMessage], scanned_tool_calls: tuple[ScannedToolCall, ...], pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], - returned_tool_calls: object, + returned_tool_calls: Sequence[object] | None, ) -> None: post_guardrail_tool_calls: Final = _tool_call_shapes( returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) else tuple(item.tool_call for item in scanned_tool_calls) ) for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): From dd173a0b1b7099255c51812edb327dffddb121fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:35 -0700 Subject: [PATCH 21/46] fix(guardrails): count the PANW latest-user scan over the hoisted system prompt --- .../panw_prisma_airs/panw_prisma_airs.py | 11 +--- .../guardrail_hooks/test_panw_prisma_airs.py | 55 ++++++------------- 2 files changed, 20 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3bc0dfabefc..9002e2aea07 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): Args: texts: Flattened text entries from the framework. - messages: Original request messages (request_data["messages"]), - NOT structured_messages (which may have injected system content). + messages: The structured messages the framework flattened into ``texts``, + hoisted top-level system prompt included, so positions line up. Returns a set of scannable indices, or None on count mismatch or no user/developer message (safety fallback to existing role-filter behavior). @@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if structured_messages: # For Anthropic /v1/messages: default to latest-user-only scanning. - # Uses request_data["messages"] (original format), NOT structured_messages - # (which has injected system content from adapter translation). if self._use_latest_user_only(request_data, logging_obj): - original_messages: Final = request_data.get("messages") - if original_messages: - scannable_indices = self._get_latest_user_text_indices(texts, original_messages) + scannable_indices = self._get_latest_user_text_indices(texts, structured_messages) # Fall through to existing role filtering if: # - not Anthropic, OR flag explicitly False, OR - # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 3d7c6e06d94..f25727ebd9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -4620,46 +4620,27 @@ class TestPanwAirsLatestRoleMessageOnly: @pytest.mark.asyncio async def test_anthropic_system_plus_multiturn_no_fallback(self): - """Anthropic with top-level system + multi-turn messages[] - — latest-user works, no scan-all fallback. + """Anthropic with a top-level system prompt and multi-turn messages[] + scans only the latest user turn, with no scan-all fallback. - Key scenario: Anthropic top-level `system` field causes - structured_messages to have an injected system entry, but - request_data["messages"] does NOT include it. + The Anthropic handler hoists the top-level `system` field into both + `texts` and `structured_messages`, so the latest-user walk has to + count the same entries the framework flattened. """ - handler = PanwPrismaAirsHandler( - guardrail_name="test_panw_airs", - api_key="test_api_key", - profile_name="test_profile", - default_on=True, + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, ) - # Original Anthropic messages (no system in messages array) - original_messages = [ - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - # texts extracted from original_messages (3 text entries) - texts = ["First user turn", "First assistant turn", "Latest user turn"] - - # structured_messages has an INJECTED system message from translation - structured_messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - inputs: GenericGuardrailAPIInputs = { - "texts": texts, - "structured_messages": structured_messages, - } + handler = make_handler() request_data = { "litellm_call_id": "test-call-id", "model": "anthropic/claude-sonnet-4-20250514", - "messages": original_messages, + "system": "You are a helpful assistant.", + "messages": [ + {"role": "user", "content": "First user turn"}, + {"role": "assistant", "content": "First assistant turn"}, + {"role": "user", "content": "Latest user turn"}, + ], "proxy_server_request": { "url": "http://localhost:4000/v1/messages", }, @@ -4670,13 +4651,11 @@ class TestPanwAirsLatestRoleMessageOnly: ) as mock_api: mock_api.return_value = {"action": "allow", "category": "benign"} - await handler.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", + await AnthropicMessagesHandler().process_input_messages( + data=request_data, + guardrail_to_apply=handler, ) - # Should scan ONLY the latest user message, not fall back to scan-all assert mock_api.call_count == 1 assert mock_api.call_args.kwargs["content"] == "Latest user turn" From 6264bd84bf1e0b33865028acdf53f5d64adb1e98 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:36 -0700 Subject: [PATCH 22/46] chore(ui): regenerate dashboard API types --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 8d5d5d58245d5765f8b5fd13d234d952a97e2030 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:03:55 -0700 Subject: [PATCH 23/46] fix(bedrock): grant rerank, retrieve, agent, and agentcore actions in the web identity session policy --- litellm/llms/bedrock/base_aws_llm.py | 121 +++++++++++------- .../test_web_identity_session_policy.py | 57 ++++++++- 2 files changed, 124 insertions(+), 54 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f52c1cec6a8..70869f69e3a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from functools import partial from threading import Lock +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx @@ -96,6 +97,76 @@ def _assume_role_params( ) +_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + + +class _SecureTransportCondition(TypedDict): + Bool: ReadOnly[_SecureTransportBool] + + +class _SessionPolicyStatement(TypedDict): + Sid: ReadOnly[str] + Effect: ReadOnly[Literal["Allow"]] + Action: ReadOnly[tuple[str, ...]] + Resource: ReadOnly[Literal["*"]] + Condition: ReadOnly[_SecureTransportCondition] + + +class WebIdentitySessionPolicy(TypedDict): + Version: ReadOnly[Literal["2012-10-17"]] + Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]] + + +_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "BedrockLiteLLM": ( + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + "bedrock:CountTokens", + "bedrock:Rerank", + "bedrock:Retrieve", + "bedrock:ListKnowledgeBases", + "bedrock:InvokeAgent", + "bedrock:ApplyGuardrail", + "bedrock:GetGuardrail", + "bedrock:ListGuardrails", + ), + "BedrockAgentCoreLiteLLM": ( + "bedrock-agentcore:InvokeAgentRuntime", + "bedrock-agentcore:InvokeGateway", + ), + "ClaudePlatformLiteLLM": ( + "aws-external-anthropic:CreateInference", + "aws-external-anthropic:CreateBatchInference", + "aws-external-anthropic:CancelBatchInference", + "aws-external-anthropic:DeleteBatchInference", + "aws-external-anthropic:CountTokens", + "aws-external-anthropic:Get*", + "aws-external-anthropic:List*", + ), + "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), + } +) + +_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) + + +def build_web_identity_session_policy() -> WebIdentitySessionPolicy: + return WebIdentitySessionPolicy( + Version="2012-10-17", + Statement=tuple( + _SessionPolicyStatement( + Sid=sid, + Effect="Allow", + Action=actions, + Resource="*", + Condition=_SECURE_TRANSPORT_ONLY, + ) + for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() + ), + ) + + class BedrockRequestTarget(BaseModel): aws_region_name: str aws_bedrock_runtime_endpoint: str | None @@ -940,60 +1011,12 @@ class BaseAWSLLM(SignsRequestsWithAWS): # auth only (static creds + IRSA take other code paths). # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html - bedrock_session_policy: Final = { - "Version": "2012-10-17", - "Statement": [ - { - "Sid": "BedrockLiteLLM", - "Effect": "Allow", - "Action": [ - "bedrock:InvokeModel", - "bedrock:InvokeModelWithResponseStream", - "bedrock:CountTokens", - "bedrock:ApplyGuardrail", - "bedrock:GetGuardrail", - "bedrock:ListGuardrails", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - # Claude Platform on AWS (added by #27678 for the - # ``bedrock/claude_platform/`` route) lives under - # a separate IAM action namespace; without these entries - # the OIDC path 403s on every claude_platform request - # even with a fully permissive identity policy (#30200). - { - "Sid": "ClaudePlatformLiteLLM", - "Effect": "Allow", - "Action": [ - "aws-external-anthropic:CreateInference", - "aws-external-anthropic:CreateBatchInference", - "aws-external-anthropic:CancelBatchInference", - "aws-external-anthropic:DeleteBatchInference", - "aws-external-anthropic:CountTokens", - "aws-external-anthropic:Get*", - "aws-external-anthropic:List*", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - { - "Sid": "BedrockMantleLiteLLM", - "Effect": "Allow", - "Action": [ - "bedrock-mantle:CreateInference", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - ], - } assume_role_params: Final = { "RoleArn": aws_role_name, "RoleSessionName": aws_session_name, "WebIdentityToken": oidc_token, "DurationSeconds": 3600, - "Policy": json.dumps(bedrock_session_policy, separators=(",", ":")), + "Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")), } # Add ExternalId parameter if provided diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 3ea840519f9..cd5b14c8d37 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -32,6 +32,8 @@ action. import base64 import json from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -49,9 +51,9 @@ _CLAUDE_PLATFORM_ACTIONS = { } -def _captured_policy() -> dict: - """Run _auth_with_web_identity_token under mocks + return the parsed - Policy dict that was actually sent to STS.""" +def _captured_policy_document() -> str: + """Run _auth_with_web_identity_token under mocks + return the Policy + JSON document that was actually sent to STS.""" from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM base = BaseAWSLLM() @@ -84,8 +86,15 @@ def _captured_policy() -> dict: mock_sts.assume_role_with_web_identity.assert_called_once() kwargs = mock_sts.assume_role_with_web_identity.call_args.kwargs - policy_str = kwargs["Policy"] - return json.loads(policy_str) + return kwargs["Policy"] + + +def _captured_policy() -> dict: + return json.loads(_captured_policy_document()) + + +def _granted_actions(policy: dict) -> frozenset[str]: + return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"]) def _statement_by_sid(policy: dict, sid: str) -> dict: @@ -308,3 +317,41 @@ class TestPolicyTransportConditions: "ClaudePlatformLiteLLM must require aws:SecureTransport=true " "to keep parity with the bedrock statement" ) + + +_STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048 + +_BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType( + { + "model/{model_id}/invoke": "bedrock:InvokeModel", + "model/{model_id}/invoke-with-response-stream": "bedrock:InvokeModelWithResponseStream", + "model/{model_id}/converse": "bedrock:InvokeModel", + "model/{model_id}/converse-stream": "bedrock:InvokeModelWithResponseStream", + "model/{model_id}/count-tokens": "bedrock:CountTokens", + "guardrail/{guardrail_id}/version/{version}/apply": "bedrock:ApplyGuardrail", + "rerank": "bedrock:Rerank", + "knowledgebases/{knowledge_base_id}/retrieve": "bedrock:Retrieve", + "knowledgebases": "bedrock:ListKnowledgeBases", + "agents/{agent_id}/agentAliases/{alias_id}/sessions/{session_id}/text": "bedrock:InvokeAgent", + "runtimes/{agent_runtime_arn}/invocations": "bedrock-agentcore:InvokeAgentRuntime", + "mcp": "bedrock-agentcore:InvokeGateway", + } +) + + +class TestSessionPolicyGrantsEveryBedrockRoute: + """LIT-7348: ``/rerank`` authorizes against ``bedrock:Rerank``, which the + ceiling never granted, so rerank 403d on web identity auth while static + credentials and IRSA worked. Each route the bedrock package signs with the + web identity session maps to the IAM action it authorizes against, and the + ceiling must grant every one of them.""" + + @pytest.mark.parametrize(("route", "action"), sorted(_BEDROCK_ROUTE_ACTIONS.items())) + def test_route_action_is_granted_by_the_ceiling(self, route: str, action: str): + assert action in _granted_actions(_captured_policy()), ( + f"/{route} authorizes against {action}, which the session policy does not grant, " + "so it 403s on web identity auth" + ) + + def test_policy_document_fits_the_sts_plaintext_limit(self): + assert len(_captured_policy_document()) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT From 7255a201a341b6e15f68a5a02a7abb0d3f6ba759 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:19:15 +0000 Subject: [PATCH 24/46] fix(health): resolve stored credentials for realtime checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 38 +++++++++++++------- tests/test_litellm/realtime_api/test_main.py | 32 +++++++++++++++++ 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 44c47af57f4..59aa63eb8a5 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -28,7 +28,7 @@ from litellm.types.realtime import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders -from litellm.utils import ProviderConfigManager +from litellm.utils import ProviderConfigManager, load_credentials_from_list from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -56,6 +56,12 @@ base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) +def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: + hydrated: Final = dict(model_params) + load_credentials_from_list(hydrated) + return MappingProxyType(hydrated) + + def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: if "model" not in session: return session @@ -629,34 +635,40 @@ async def _realtime_health_check( """ import websockets + resolved_params: Final = _model_params_with_stored_credentials(model_params or _EMPTY_MODEL_PARAMS) + resolved_api_key: Final = cast(str | None, api_key or resolved_params.get("api_key")) + resolved_api_base: Final = cast(str | None, api_base or resolved_params.get("api_base")) + resolved_api_version: Final = cast(str | None, api_version or resolved_params.get("api_version")) url: str | None = None auth_headers: Final = _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, - api_key=api_key, - model_params=model_params or _EMPTY_MODEL_PARAMS, + api_key=resolved_api_key, + model_params=resolved_params, ) if custom_llm_provider == "azure": resolved_protocol, azure_query_params = _azure_realtime_health_protocol( model=model, realtime_protocol=realtime_protocol, - model_params=model_params or _EMPTY_MODEL_PARAMS, + model_params=resolved_params, ) url = azure_realtime._construct_url( - api_base=api_base or "", + api_base=resolved_api_base or "", model=model, - api_version=api_version or "2024-10-01-preview", + api_version=resolved_api_version or "2024-10-01-preview", realtime_protocol=resolved_protocol, query_params=azure_query_params, ) elif custom_llm_provider == "openai": url = openai_realtime._construct_url( - api_base=api_base or "https://api.openai.com/", + api_base=resolved_api_base or "https://api.openai.com/", query_params={"model": model}, ) elif custom_llm_provider == "xai": - url = xai_realtime._construct_url(api_base=api_base or "https://api.x.ai/v1", query_params={"model": model}) + url = xai_realtime._construct_url( + api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model} + ) elif custom_llm_provider == "vertex_ai": - vertex_model_params: Final = model_params or {} + vertex_model_params: Final = dict(resolved_params) resolved_location: Final = vertex_llm_base.get_vertex_region( vertex_region=VertexBase.safe_get_vertex_ai_location(vertex_model_params), model=model, @@ -675,19 +687,19 @@ async def _realtime_health_check( project=resolved_project, location=resolved_location, ) - url = vertex_realtime_config.get_complete_url(api_base=api_base, model=model) - ssl_context = get_shared_realtime_ssl_context() + url = vertex_realtime_config.get_complete_url(api_base=resolved_api_base, model=model) + vertex_ssl_context: Final = get_shared_realtime_ssl_context() headers: Final = vertex_realtime_config.validate_environment(headers={}, model=model, api_key=None) async with websockets.connect( url, additional_headers=headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, - ssl=ssl_context, + ssl=vertex_ssl_context, ): return True else: raise ValueError(f"Unsupported model: {model}") - ssl_context = get_shared_realtime_ssl_context() + ssl_context: Final = get_shared_realtime_ssl_context() async with websockets.connect( url, additional_headers=auth_headers, diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index d3d41c5b54b..25eeb0c407f 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm +from litellm.models.credentials import CredentialItem from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model @@ -224,9 +225,11 @@ def test_client_secret_forwards_nested_transcription_model_untouched(monkeypatch class _CapturingConnect: def __init__(self) -> None: self.url: str | None = None + self.kwargs: dict[str, object] = {} def __call__(self, url: str, **kwargs: object) -> "_CapturingConnect": self.url = url + self.kwargs = kwargs return self async def __aenter__(self) -> MagicMock: @@ -241,6 +244,35 @@ class _CapturingConnect: return None +@pytest.mark.asyncio +async def test_azure_health_check_resolves_stored_credentials(monkeypatch): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="azure-rt", + credential_values={ + "api_key": "sk-from-credential", + "api_base": "https://example.openai.azure.com", + "api_version": "2025-04-01-preview", + }, + credential_info={}, + ) + ], + ) + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model="gpt-realtime", + custom_llm_provider="azure", + api_key=None, + model_params={"model": "azure/gpt-realtime", "litellm_credential_name": "azure-rt"}, + ) + assert connect.kwargs["additional_headers"] == {"api-key": "sk-from-credential"} + assert connect.url is not None and connect.url.startswith("wss://example.openai.azure.com") + + @pytest.mark.asyncio async def test_azure_health_check_probes_ga_transcription_url_for_transcription_model(local_model_cost_map): """Regression for LIT-6240: transcription-only models (mode audio_transcription From 0a0a9509dcf75c55c138c1d2b6e8f0219f089523 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:21:23 -0700 Subject: [PATCH 25/46] test(bedrock): type the session policy test helpers with the policy TypedDict --- .../test_web_identity_session_policy.py | 36 +++++++------------ 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index cd5b14c8d37..2e84c3ca36a 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -37,6 +37,9 @@ from typing import Final from unittest.mock import MagicMock, patch import pytest +from pydantic import TypeAdapter + +from litellm.llms.bedrock.base_aws_llm import WebIdentitySessionPolicy, _SessionPolicyStatement # Actions the Claude Platform on AWS service is documented to call. # Source: AWS IAM action reference + the #27678 surface area. @@ -89,15 +92,18 @@ def _captured_policy_document() -> str: return kwargs["Policy"] -def _captured_policy() -> dict: - return json.loads(_captured_policy_document()) +_SESSION_POLICY_ADAPTER: Final = TypeAdapter(WebIdentitySessionPolicy) -def _granted_actions(policy: dict) -> frozenset[str]: +def _captured_policy() -> WebIdentitySessionPolicy: + return _SESSION_POLICY_ADAPTER.validate_python(json.loads(_captured_policy_document())) + + +def _granted_actions(policy: WebIdentitySessionPolicy) -> frozenset[str]: return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"]) -def _statement_by_sid(policy: dict, sid: str) -> dict: +def _statement_by_sid(policy: WebIdentitySessionPolicy, sid: str) -> _SessionPolicyStatement: for stmt in policy["Statement"]: if stmt.get("Sid") == sid: return stmt @@ -111,7 +117,6 @@ class TestWebIdentitySessionPolicyShape: def test_policy_parses_as_valid_iam_document(self): policy = _captured_policy() assert policy["Version"] == "2012-10-17" - assert isinstance(policy["Statement"], list) assert len(policy["Statement"]) >= 2 def test_bedrock_statement_actions_preserved(self): @@ -146,16 +151,7 @@ class TestClaudePlatformActionsCovered: @pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS)) def test_claude_platform_action_present(self, action: str): - policy = _captured_policy() - # Action may live in any Statement — search across all. - all_actions: set = set() - for stmt in policy["Statement"]: - stmt_actions = stmt.get("Action") - if isinstance(stmt_actions, str): - all_actions.add(stmt_actions) - elif isinstance(stmt_actions, list): - all_actions.update(stmt_actions) - assert action in all_actions, ( + assert action in _granted_actions(_captured_policy()), ( f"{action} missing from session policy — " f"bedrock/claude_platform/* requests will 403 on OIDC auth" ) @@ -188,15 +184,7 @@ class TestBedrockMantleActionsCovered: action" even when the role's identity policy grants it.""" def test_bedrock_mantle_create_inference_present(self): - policy = _captured_policy() - all_actions: set = set() - for stmt in policy["Statement"]: - stmt_actions = stmt.get("Action") - if isinstance(stmt_actions, str): - all_actions.add(stmt_actions) - elif isinstance(stmt_actions, list): - all_actions.update(stmt_actions) - assert "bedrock-mantle:CreateInference" in all_actions, ( + assert "bedrock-mantle:CreateInference" in _granted_actions(_captured_policy()), ( "bedrock-mantle:CreateInference missing from session policy — " "bedrock_mantle/* requests will 403 on OIDC/WIF auth" ) From 123b76e2d80f614a7e3595d74ac7e89d2d7a1335 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:34:49 +0000 Subject: [PATCH 26/46] fix(health): annotate credential resolution casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 59aa63eb8a5..336fb3cef12 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -636,9 +636,15 @@ async def _realtime_health_check( import websockets resolved_params: Final = _model_params_with_stored_credentials(model_params or _EMPTY_MODEL_PARAMS) - resolved_api_key: Final = cast(str | None, api_key or resolved_params.get("api_key")) - resolved_api_base: Final = cast(str | None, api_base or resolved_params.get("api_base")) - resolved_api_version: Final = cast(str | None, api_version or resolved_params.get("api_version")) + resolved_api_key: Final = cast( # cast-ok: provider parameters expose optional string credentials + str | None, api_key or resolved_params.get("api_key") + ) + resolved_api_base: Final = cast( # cast-ok: provider parameters expose optional string endpoints + str | None, api_base or resolved_params.get("api_base") + ) + resolved_api_version: Final = cast( # cast-ok: provider parameters expose optional string versions + str | None, api_version or resolved_params.get("api_version") + ) url: str | None = None auth_headers: Final = _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, From e668a313847c9fd946d9a84c7d42e1616e9b9431 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:43:48 +0000 Subject: [PATCH 27/46] fix(health): keep realtime credential hydration immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 11 +++++++---- tests/test_litellm/realtime_api/test_main.py | 5 ++++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 336fb3cef12..6fdadc52f13 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -28,8 +28,9 @@ from litellm.types.realtime import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders -from litellm.utils import ProviderConfigManager, load_credentials_from_list +from litellm.utils import ProviderConfigManager +from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.common_utils import get_azure_ad_token @@ -57,9 +58,11 @@ _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: - hydrated: Final = dict(model_params) - load_credentials_from_list(hydrated) - return MappingProxyType(hydrated) + credential_name: Final = model_params.get("litellm_credential_name") + credential_values: Final = ( + CredentialAccessor.get_credential_values(credential_name) if isinstance(credential_name, str) else {} + ) + return MappingProxyType({**credential_values, **model_params}) def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 25eeb0c407f..dcd4c01730d 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -267,10 +267,13 @@ async def test_azure_health_check_resolves_stored_credentials(monkeypatch): model="gpt-realtime", custom_llm_provider="azure", api_key=None, + realtime_protocol="beta", model_params={"model": "azure/gpt-realtime", "litellm_credential_name": "azure-rt"}, ) assert connect.kwargs["additional_headers"] == {"api-key": "sk-from-credential"} - assert connect.url is not None and connect.url.startswith("wss://example.openai.azure.com") + assert connect.url is not None + assert connect.url.startswith("wss://example.openai.azure.com") + assert "api-version=2025-04-01-preview" in connect.url @pytest.mark.asyncio From ae90f1a45891debf622d3cb531163e4fb7f21399 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:28:47 -0700 Subject: [PATCH 28/46] fix(guardrails): type the request payload handed to the Anthropic write-back --- .../llms/anthropic/chat/guardrail_translation/handler.py | 7 ++++--- .../test_anthropic_guardrail_handler.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b438d168b52..b8e179e7274 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1079,14 +1079,15 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - data: dict, - responses: list[str], + data: dict[str, object], # mutable-ok: API message payload + responses: Sequence[str], scanned: tuple[ScannedText, ...], ) -> None: """ Apply guardrail responses back to the top-level system prompt and the input messages. """ - messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () + raw_messages: Final = data.get("messages") + messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else () for item, guardrail_response in zip(scanned, responses): match item.target: case SystemStringTarget(): diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 4b47b7d8c4a..51c3751dcf8 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2200,7 +2200,7 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: inputs, the same way the chat completions handler hands over system messages and tool_calls.""" @staticmethod - def _tool_use_conversation(system): + def _tool_use_conversation(system: str) -> dict[str, Any]: return { "model": "claude-sonnet-4-5", "system": system, From 1111658e16cea78af5509d4dd9a9ca4b0e610b9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:31:45 -0700 Subject: [PATCH 29/46] fix(bedrock): grant the AgentCore for-user invoke action in the web identity session policy The chat and A2A AgentCore handlers send X-Amzn-Bedrock-AgentCore-Runtime-User-Id when runtimeUserId is set, and AWS requires bedrock-agentcore:InvokeAgentRuntimeForUser alongside InvokeAgentRuntime on that call, so the ceiling now carries both. The role identity policy still decides whether a given role may use it The invalid-token test now uses a neutral example audience --- litellm/llms/bedrock/base_aws_llm.py | 1 + .../llms/bedrock/test_web_identity_session_policy.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 70869f69e3a..385d5898569 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -133,6 +133,7 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map ), "BedrockAgentCoreLiteLLM": ( "bedrock-agentcore:InvokeAgentRuntime", + "bedrock-agentcore:InvokeAgentRuntimeForUser", "bedrock-agentcore:InvokeGateway", ), "ClaudePlatformLiteLLM": ( diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 2e84c3ca36a..cbb69b4ceed 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -230,7 +230,7 @@ class TestInvalidIdentityTokenSurfacesAudience: operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a prod instance.""" - _AUD = "https://guidepoint.litellm-prod.ai" + _AUD = "https://gateway.example.com" _ISS = "https://accounts.google.com" _STS_MESSAGE = ( "An error occurred (InvalidIdentityToken) when calling the " @@ -322,6 +322,9 @@ _BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType( "knowledgebases": "bedrock:ListKnowledgeBases", "agents/{agent_id}/agentAliases/{alias_id}/sessions/{session_id}/text": "bedrock:InvokeAgent", "runtimes/{agent_runtime_arn}/invocations": "bedrock-agentcore:InvokeAgentRuntime", + "runtimes/{agent_runtime_arn}/invocations with X-Amzn-Bedrock-AgentCore-Runtime-User-Id": ( + "bedrock-agentcore:InvokeAgentRuntimeForUser" + ), "mcp": "bedrock-agentcore:InvokeGateway", } ) From bebc76316cd8269c20a5682147745f16bbd5a568 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 21:42:57 -0700 Subject: [PATCH 30/46] fix(cli): label savings cost bars with the auto-router name --- .../proxy/client/cli/commands/statusline_script.py | 5 ++--- .../proxy/client/cli/test_statusline_script.py | 12 +++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 5493586f627..0e1c1b25e0f 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",) CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",) CODEX_STOP_EVENT: Final = "Stop" SYNTHETIC_MODEL: Final = "" -LITELLM_LABEL: Final = "LiteLLM" RESET: Final = "\033[0m" BOLD: Final = "\033[1m" DIM: Final = "\033[90m" @@ -316,9 +315,9 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") peak: Final = max(session.spend, session.baseline_spend) - label_width: Final = max(len(LITELLM_LABEL), len(reference)) + label_width: Final = max(len(session.router_name), len(reference)) rows: Final = ( - (LITELLM_LABEL, session.spend, LITELLM_COLOR), + (session.router_name, session.spend, LITELLM_COLOR), (reference, session.baseline_spend, BASELINE_COLOR), ) lines: Final = ( diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 5c0cf6b5703..2b812932542 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -253,10 +253,18 @@ class TestRender: text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) assert text.splitlines() == [ "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", - "LiteLLM ████░░░░░░ $0.14", + "claude-auto ████░░░░░░ $0.14", "Claude Opus 5 ██████████ $0.38", ] + def test_a_long_router_name_keeps_both_cost_bars_aligned(self, config_dir: Path) -> None: + session: Final = RECORDED._replace(router_name="engineering-smart-router") + text: Final = render("claude-sonnet-5", session, config_dir, use_color=False, bar_width=10) + assert text.splitlines()[1:] == [ + "engineering-smart-router ████░░░░░░ $0.14", + "Claude Opus 5 ██████████ $0.38", + ] + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, @@ -312,6 +320,7 @@ class TestClaudeCodeMode: text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert text.splitlines()[1].startswith("claude-auto ") def test_a_discovered_display_name_labels_the_sessions_model( self, tmp_path: Path, transcript: Path, config_dir: Path @@ -379,6 +388,7 @@ class TestCodexMode: out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) message = json.loads(out)["systemMessage"] assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.splitlines()[2].startswith("claude-auto ") assert message.startswith("\n") assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] From 31b34f7767cf8426293e98723f8e80cc667a1cab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:46:02 -0700 Subject: [PATCH 31/46] test(guardrails): type the Anthropic write-back test helper --- .../test_anthropic_guardrail_handler.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 51c3751dcf8..32786fc5057 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, @@ -2163,14 +2164,14 @@ class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): super().__init__() self.return_copies = return_copies self.replacement_arguments = replacement_arguments - self.seen_tool_calls: list[dict] = [] + self.seen_tool_calls: list[dict[str, object]] = [] async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: Optional[LiteLLMLoggingObj] = None, ) -> GenericGuardrailAPIInputs: outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) tool_calls = list(outputs.get("tool_calls") or []) From 7e3d7178b4194e642fe4a9a3a5ffae79945f8612 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 21:52:01 -0700 Subject: [PATCH 32/46] test(spend): reconcile concurrent requests and daily activity --- .../spend_tracking/spend_reconciliation.py | 109 +++++++++++++ .../spend_tracking/test_spend_tracking_e2e.py | 56 ++----- .../test_team_daily_activity_e2e.py | 144 +++++++++++++++--- 3 files changed, 244 insertions(+), 65 deletions(-) create mode 100644 tests/e2e/quota_management/spend_tracking/spend_reconciliation.py diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py new file mode 100644 index 00000000000..8fcfea3f296 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -0,0 +1,109 @@ +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from math import isclose +from typing import Final + +from e2e_config import provider_edge_base, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody +from spend_e2e_client import SpendClient + +INPUT_RATE: Final = 0.00004 +OUTPUT_RATE: Final = 0.00008 + + +@dataclass(frozen=True) +class TeamTraffic: + team_id: str + key: str + responses: tuple[ChatResponse, ...] + + @property + def prompt_tokens(self) -> int: + return sum(response.usage.prompt_tokens or 0 for response in self.responses if response.usage) + + @property + def completion_tokens(self) -> int: + return sum(response.usage.completion_tokens or 0 for response in self.responses if response.usage) + + @property + def spend(self) -> float: + return self.prompt_tokens * INPUT_RATE + self.completion_tokens * OUTPUT_RATE + + +def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[TeamTraffic, ...]: + base: Final = provider_edge_base("openai") + model: Final = f"e2e-reconciliation-{unique_marker()}" + model_id: Final = client.proxy.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6-luna", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + def team_traffic() -> TeamTraffic: + team: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-{unique_marker()}")) + resources.defer(lambda: client.proxy.delete_team(team)) + key: Final = client.proxy.generate_key(KeyGenerateBody(team_id=team, models=[model])) + resources.defer(lambda: client.proxy.delete_key(key)) + + prompts: Final = tuple(f"Reply with one word. {index} {unique_marker()}" for index in range(7)) + + def call(index: int) -> ChatResponse: + response: Final = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=prompts[index])], + max_completion_tokens=128, + ), + ) + ) + assert response.id, "successful response must have an ID" + assert response.usage is not None, "successful response must have usage" + assert response.usage.prompt_tokens is not None and response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens is not None and response.usage.completion_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + assert not response.usage.cache_creation_input_tokens + assert not response.usage.cache_read_input_tokens + assert not response.usage.prompt_tokens_details or not response.usage.prompt_tokens_details.cached_tokens + return response + + sequential: Final = call(0) + with ThreadPoolExecutor(max_workers=6) as pool: + concurrent: Final = tuple(pool.map(call, range(1, 7))) + return TeamTraffic(team, key, (sequential, *concurrent)) + + return tuple(team_traffic() for _ in range(2)) + + +def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: + expected_ids: Final = frozenset(response.id for response in traffic.responses) + assert len(expected_ids) == len(traffic.responses), "responses must have distinct IDs" + rows: Final = client.poll_logs_for_key( + traffic.key, + min_rows=len(traffic.responses), + predicate=lambda values: frozenset(row.request_id for row in values) == expected_ids, + ) + assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs" + assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response" + by_id: Final = {row.request_id: row for row in rows} + for response in traffic.responses: + row = by_id[response.id] + usage = response.usage + assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None + assert row.team_id == traffic.team_id + assert row.status == "success" + assert row.cache_hit != "True" + assert row.prompt_tokens == usage.prompt_tokens + assert row.completion_tokens == usage.completion_tokens + assert row.total_tokens == usage.total_tokens + expected_cost = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index c5d76d44580..750285cbfb6 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -17,13 +17,12 @@ fails the test; a pricing or token-count drift does not. import time from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor +from math import isclose import pytest - -from e2e_http import Result, Success +from e2e_http import Success from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, SpendLogs, SpendLogsParams +from models import LiteLLMParamsBody, SpendLogs, SpendLogsParams from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap pytestmark = pytest.mark.e2e @@ -280,51 +279,18 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N ), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" +@pytest.mark.replayable @pytest.mark.covers("quota_management.spend_tracking.concurrent_burst.loses_no_spend") def test_burst_of_concurrent_calls_loses_no_spend( - client: SpendClient, scoped_key: str + client: SpendClient, resources: ResourceManager ) -> None: - """Six concurrent calls on one key: every call lands its own spend row under a - distinct request_id and the key aggregate equals the sum of the rows. - Sequential accuracy is covered by test_key_spend_equals_sum_of_logs; this pins - the concurrent increment path (parallel writers racing on one key's counter), - where a lost update can never be reproduced by sequential calls.""" - burst = 6 + from spend_reconciliation import assert_logs_match, create_traffic - def call(idx: int) -> Result[ChatResponse]: - return client.chat( - scoped_key, - "gemini-2.5-flash", - f"burst call {idx} {unique_marker()}", - max_tokens=16, - ) - - with ThreadPoolExecutor(max_workers=burst) as pool: - results = tuple(pool.map(call, range(burst))) - failed = [r for r in results if not is_ok(r)] - assert not failed, f"{len(failed)}/{burst} burst calls failed; first: {failed[0]}" - - rows = client.poll_logs_for_key( - scoped_key, - min_rows=burst, - predicate=lambda rs: len([r for r in rs if (r.spend or 0) > 0]) >= burst, - ) - costed = [r for r in rows if (r.spend or 0) > 0] - assert len(costed) >= burst, ( - f"only {len(costed)}/{burst} burst calls produced a costed row - " - f"rows lost under concurrency: {_summarize(rows)}" - ) - request_ids = [r.request_id for r in costed] - assert len(set(request_ids)) == len(request_ids), ( - f"concurrent rows collapsed onto shared request_ids: {_summarize(rows)}" - ) - - logs_total = sum((r.spend or 0) for r in rows) - key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999) - assert _approx_equal(key_spend, logs_total), ( - f"key aggregate {key_spend} != sum of {len(rows)} rows {logs_total} - " - f"spend increments lost under concurrency: {_summarize(rows)}" - ) + traffic = create_traffic(client, resources) + for team in traffic: + assert_logs_match(client, team) + key_spend = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) + assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9) @pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total") diff --git a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py index ed0a6af4ec9..dca5572f510 100644 --- a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py @@ -7,13 +7,17 @@ missing start/end dates are rejected. from __future__ import annotations +import time from datetime import datetime, timedelta, timezone +from math import isclose +from typing import Final import pytest from e2e_http import ProbeResult -from models import DateRangeParams +from lifecycle import ResourceManager from pydantic import BaseModel from spend_e2e_client import SpendClient +from spend_reconciliation import assert_logs_match, create_traffic pytestmark = pytest.mark.e2e @@ -24,22 +28,45 @@ class TeamDailyActivityParams(BaseModel): start_date: str | None = None end_date: str | None = None page: int = 1 + page_size: int = 1 + team_ids: str | None = None class TeamDailyActivityRow(BaseModel): date: str metrics: TeamDailyActivityMetrics + breakdown: TeamDailyActivityBreakdown class TeamDailyActivityMetrics(BaseModel): spend: float total_tokens: int + prompt_tokens: int + completion_tokens: int + api_requests: int + successful_requests: int + failed_requests: int + + +class TeamDailyActivityEntity(BaseModel): + metrics: TeamDailyActivityMetrics + + +class TeamDailyActivityBreakdown(BaseModel): + entities: dict[str, TeamDailyActivityEntity] class TeamDailyActivityMetadata(BaseModel): page: int total_pages: int has_more: bool + total_spend: float + total_prompt_tokens: int + total_completion_tokens: int + total_tokens: int + total_api_requests: int + total_successful_requests: int + total_failed_requests: int class TeamDailyActivityResponse(BaseModel): @@ -47,32 +74,109 @@ class TeamDailyActivityResponse(BaseModel): metadata: TeamDailyActivityMetadata -def _range_days(days: int) -> DateRangeParams: - end = datetime.now(timezone.utc).date() - start = end - timedelta(days=days) - return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) - - def _probe(client: SpendClient, params: BaseModel) -> ProbeResult: return client.proxy.transport.probe(ROUTE, params=params) class TestTeamDailyActivity: + @pytest.mark.replayable @pytest.mark.covers("mgmt.team.daily_activity.happy_path") - @pytest.mark.parametrize("days", [1, 7, 30]) - def test_valid_date_range_returns_results_and_metadata(self, client: SpendClient, days: int) -> None: - result = _probe(client, _range_days(days)) - assert result.status_code == 200, ( - f"{ROUTE} range={days}d must be 200, got {result.status_code}: {result.body[:600]}" + def test_valid_date_range_returns_results_and_metadata( + self, client: SpendClient, resources: ResourceManager + ) -> None: + started: Final = datetime.now(timezone.utc).date() + traffic: Final = create_traffic(client, resources) + for team in traffic: + assert_logs_match(client, team) + ended: Final = datetime.now(timezone.utc).date() + team_ids: Final = ",".join(team.team_id for team in traffic) + + def fetch( + page: int, start: str = started.isoformat(), end: str = ended.isoformat() + ) -> TeamDailyActivityResponse: + result: Final = _probe( + client, + TeamDailyActivityParams( + start_date=start, + end_date=end, + page=page, + page_size=1, + team_ids=team_ids, + ), + ) + assert result.status_code == 200, f"daily activity failed: {result.status_code} {result.body[:300]}" + return TeamDailyActivityResponse.model_validate_json(result.body) + + def pages() -> tuple[TeamDailyActivityResponse, ...]: + first: Final = fetch(1) + assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups" + return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1))) + + deadline: Final = time.monotonic() + client.proxy.poll_timeout + while True: + observed = pages() + if sum(page.metadata.total_api_requests for page in observed) >= sum(len(t.responses) for t in traffic): + break + if time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + + assert len(observed) >= 2, "two teams must exercise a page boundary" + for index, page in enumerate(observed, 1): + assert page.metadata.page == index + assert page.metadata.total_pages == len(observed) + assert page.metadata.has_more == (index < len(observed)) + assert len(page.results) == 1, "each fetched daily group must appear in results" + row = page.results[0] + assert started <= datetime.fromisoformat(row.date).date() <= ended + assert len(row.breakdown.entities) == 1 + assert row.metrics.total_tokens == page.metadata.total_tokens + assert row.metrics.prompt_tokens == page.metadata.total_prompt_tokens + assert row.metrics.completion_tokens == page.metadata.total_completion_tokens + assert row.metrics.api_requests == page.metadata.total_api_requests + assert row.metrics.successful_requests == page.metadata.total_successful_requests + assert row.metrics.failed_requests == page.metadata.total_failed_requests + assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9) + + entities: Final = tuple( + (team_id, entity.metrics) + for page in observed + for row in page.results + for team_id, entity in row.breakdown.entities.items() ) - parsed = TeamDailyActivityResponse.model_validate_json(result.body) - assert parsed.metadata.page == 1 - assert parsed.metadata.total_pages >= 1 - if parsed.results: - first = parsed.results[0] - assert first.date - assert first.metrics.spend >= 0 - assert first.metrics.total_tokens >= 0 + assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic) + for team in traffic: + metrics = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) + assert sum(m.api_requests for m in metrics) == len(team.responses) + assert sum(m.successful_requests for m in metrics) == len(team.responses) + assert sum(m.failed_requests for m in metrics) == 0 + assert sum(m.prompt_tokens for m in metrics) == team.prompt_tokens + assert sum(m.completion_tokens for m in metrics) == team.completion_tokens + assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens + assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9) + assert isclose( + sum(page.metadata.total_spend for page in observed), + sum(team.spend for team in traffic), + rel_tol=1e-6, + abs_tol=1e-9, + ) + assert sum(page.metadata.total_tokens for page in observed) == sum( + team.prompt_tokens + team.completion_tokens for team in traffic + ) + + empty_date: Final = (started - timedelta(days=7)).isoformat() + empty: Final = fetch(1, empty_date, empty_date) + assert empty.results == [] + assert empty.metadata.total_pages == 0 + assert empty.metadata.page == 1 + assert not empty.metadata.has_more + assert empty.metadata.total_spend == 0 + assert empty.metadata.total_tokens == 0 + assert empty.metadata.total_api_requests == 0 + assert empty.metadata.total_prompt_tokens == 0 + assert empty.metadata.total_completion_tokens == 0 + assert empty.metadata.total_successful_requests == 0 + assert empty.metadata.total_failed_requests == 0 @pytest.mark.covers("mgmt.team.daily_activity.missing_start_date_rejected") def test_missing_start_date_is_rejected(self, client: SpendClient) -> None: From 1b594fc93515dd70168325839679e1aad2df71be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:24:03 -0700 Subject: [PATCH 33/46] fix(guardrails): scan empty top-level system text blocks too The hoisted structured row keeps every text block of the top-level system prompt, empty ones included, while the scanned texts dropped the empty ones. Guardrails that count one text per slot then came back with more texts than the handler could place, so their rewrite was rejected. User text blocks were already scanned empty or not; the system prompt now matches. --- .../chat/guardrail_translation/handler.py | 2 +- .../test_anthropic_guardrail_handler.py | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 1f4e1487316..9f3c85b555a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -721,7 +721,7 @@ class AnthropicMessagesHandler(BaseTranslation): return tuple( ScannedText(text_str, SystemBlockTextTarget(block_idx)) for block_idx, block in enumerate(content) - if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) ) @staticmethod diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 1cacbb6be6b..287674afaa7 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2499,6 +2499,29 @@ class PerRowTextGuardrail(CustomGuardrail): return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "") for row in rows]} +class PerSlotTextGuardrail(CustomGuardrail): + """Answers one redacted text per text slot of every chat row it was shown, the + way a guardrail that counts slots per message does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="per-slot-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts + + rows = inputs.get("structured_messages") or [] + return { + **inputs, + "texts": [text.replace("123-45-6789", "") for row in rows for text in message_slot_texts(row)], + } + + class TestPerMessageTextWriteBack: """Texts that no longer pair one-to-one with what the handler extracted must be rejected by name instead of sliding onto the wrong messages.""" @@ -2537,6 +2560,25 @@ class TestPerMessageTextWriteBack: assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + @pytest.mark.asyncio + async def test_one_text_per_slot_over_a_system_prompt_with_an_empty_block_is_applied(self): + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Reply with exactly the SSN you were given."}, + ], + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerSlotTextGuardrail()) + + assert data["system"] == [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Reply with exactly the SSN you were given."}, + ] + assert data["messages"] == [{"role": "user", "content": "My SSN is ."}] + @pytest.mark.asyncio async def test_one_text_per_row_without_a_system_prompt_is_applied(self): data = { From 80d804d6f98b69ff280b2020df76cc3bd6aa07fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 22:26:20 -0700 Subject: [PATCH 34/46] test(spend): preserve multi-day coverage and immutable assertions --- .../spend_tracking/spend_reconciliation.py | 12 +++-- .../spend_tracking/test_spend_tracking_e2e.py | 13 +++-- .../test_team_daily_activity_e2e.py | 48 +++++++++++++------ 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py index 8fcfea3f296..26809874aed 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -95,9 +95,10 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs" assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response" by_id: Final = {row.request_id: row for row in rows} - for response in traffic.responses: - row = by_id[response.id] - usage = response.usage + + def assert_response(response: ChatResponse) -> None: + row: Final = by_id[response.id] + usage: Final = response.usage assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None assert row.team_id == traffic.team_id assert row.status == "success" @@ -105,5 +106,8 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert row.prompt_tokens == usage.prompt_tokens assert row.completion_tokens == usage.completion_tokens assert row.total_tokens == usage.total_tokens - expected_cost = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + expected_cost: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9) + + for response in traffic.responses: + assert_response(response) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 750285cbfb6..8a91e53e7d7 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -18,6 +18,7 @@ fails the test; a pricing or token-count drift does not. import time from collections.abc import Callable from math import isclose +from typing import Final import pytest from e2e_http import Success @@ -284,14 +285,18 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N def test_burst_of_concurrent_calls_loses_no_spend( client: SpendClient, resources: ResourceManager ) -> None: - from spend_reconciliation import assert_logs_match, create_traffic + from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic - traffic = create_traffic(client, resources) - for team in traffic: + traffic: Final = create_traffic(client, resources) + + def assert_team(team: TeamTraffic) -> None: assert_logs_match(client, team) - key_spend = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) + key_spend: Final = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9) + for team in traffic: + assert_team(team) + @pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total") def test_spend_logs_v2_pagination_caps_pages_and_keeps_total( diff --git a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py index dca5572f510..ef635e59743 100644 --- a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py @@ -15,9 +15,10 @@ from typing import Final import pytest from e2e_http import ProbeResult from lifecycle import ResourceManager +from proxy_client import Converged, await_converged from pydantic import BaseModel from spend_e2e_client import SpendClient -from spend_reconciliation import assert_logs_match, create_traffic +from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic pytestmark = pytest.mark.e2e @@ -92,7 +93,7 @@ class TestTeamDailyActivity: team_ids: Final = ",".join(team.team_id for team in traffic) def fetch( - page: int, start: str = started.isoformat(), end: str = ended.isoformat() + page: int, start: str = (started - timedelta(days=1)).isoformat(), end: str = ended.isoformat() ) -> TeamDailyActivityResponse: result: Final = _probe( client, @@ -112,22 +113,27 @@ class TestTeamDailyActivity: assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups" return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1))) - deadline: Final = time.monotonic() + client.proxy.poll_timeout - while True: - observed = pages() - if sum(page.metadata.total_api_requests for page in observed) >= sum(len(t.responses) for t in traffic): - break - if time.monotonic() >= deadline: - break - time.sleep(client.proxy.poll_interval) + outcome: Final = await_converged( + pages, + converged=lambda values: ( + sum(page.metadata.total_api_requests for page in values) >= sum(len(team.responses) for team in traffic) + ), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + observed: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result + assert observed is not None, "daily aggregation must return a response before the deadline" assert len(observed) >= 2, "two teams must exercise a page boundary" - for index, page in enumerate(observed, 1): + + def assert_page(index: int, page: TeamDailyActivityResponse) -> None: assert page.metadata.page == index assert page.metadata.total_pages == len(observed) assert page.metadata.has_more == (index < len(observed)) assert len(page.results) == 1, "each fetched daily group must appear in results" - row = page.results[0] + row: Final = page.results[0] assert started <= datetime.fromisoformat(row.date).date() <= ended assert len(row.breakdown.entities) == 1 assert row.metrics.total_tokens == page.metadata.total_tokens @@ -138,6 +144,9 @@ class TestTeamDailyActivity: assert row.metrics.failed_requests == page.metadata.total_failed_requests assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9) + for index, page in enumerate(observed, 1): + assert_page(index, page) + entities: Final = tuple( (team_id, entity.metrics) for page in observed @@ -145,8 +154,9 @@ class TestTeamDailyActivity: for team_id, entity in row.breakdown.entities.items() ) assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic) - for team in traffic: - metrics = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) + + def assert_team(team: TeamTraffic) -> None: + metrics: Final = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) assert sum(m.api_requests for m in metrics) == len(team.responses) assert sum(m.successful_requests for m in metrics) == len(team.responses) assert sum(m.failed_requests for m in metrics) == 0 @@ -154,6 +164,10 @@ class TestTeamDailyActivity: assert sum(m.completion_tokens for m in metrics) == team.completion_tokens assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9) + + for team in traffic: + assert_team(team) + assert isclose( sum(page.metadata.total_spend for page in observed), sum(team.spend for team in traffic), @@ -164,6 +178,12 @@ class TestTeamDailyActivity: team.prompt_tokens + team.completion_tokens for team in traffic ) + for days in (7, 30): + assert ( + tuple(fetch(page, (started - timedelta(days=days)).isoformat()) for page in range(1, len(observed) + 1)) + == observed + ), f"{days}-day activity must preserve the same isolated groups and totals" + empty_date: Final = (started - timedelta(days=7)).isoformat() empty: Final = fetch(1, empty_date, empty_date) assert empty.results == [] From 6764ab2673942e7a32a3a12414c5e7ecec64a8fa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:29:32 -0700 Subject: [PATCH 35/46] test(router): assert num_retries_per_request as a per-group cap that resets per fallback hop #40930 (LIT-7505) changed num_retries_per_request from a request-wide cap to a per-model-group cap that resets on every fallback hop, and its own comment in litellm/__init__.py names that contract. The legacy test_async_fallbacks_max_retries_per_request still asserted the old request-wide reading (previous_models == 0), so the CircleCI router suite has been red on main since that merge for every run-ci PR. The test now reads the flat RetryAttemptRecord entries the fallback call carries and asserts the new contract directly: every record is from the first group, the retry at attempted_retries 0 is the real AuthenticationError, and each later attempt was refused with "Max retries per request hit!". --- tests/local_testing/test_router_fallbacks.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 82b832f89fd..5d7955ad34a 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -5,6 +5,7 @@ import asyncio import os import time import traceback +from typing import Final import pytest @@ -13,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RetryAttemptRecord from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -21,15 +23,15 @@ class MyCustomHandler(CustomLogger): success: bool = False failure: bool = False previous_models: int = 0 + previous_model_records: tuple[RetryAttemptRecord, ...] = () def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call") print( f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}" ) - self.previous_models = len( - kwargs["litellm_params"]["metadata"].get("previous_models", []) - ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} + self.previous_model_records = tuple(kwargs["litellm_params"]["metadata"].get("previous_models", ())) + self.previous_models = len(self.previous_model_records) print(f"self.previous_models: {self.previous_models}") def log_post_api_call(self, kwargs, response_obj, start_time, end_time): @@ -718,7 +720,14 @@ async def test_async_fallbacks_max_retries_per_request(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models == 0 # 0 retries, 0 fallback + records: Final = customHandler.previous_model_records + assert customHandler.previous_models == len(records) + assert records + assert {record["model_group"] for record in records} == {"azure/gpt-3.5-turbo"} + assert next(record["exception_type"] for record in records if record["attempted_retries"] == 0) == "AuthenticationError" + refused_retries: Final = tuple(record for record in records if record["attempted_retries"]) + assert refused_retries + assert all("Max retries per request hit!" in record["exception_string"] for record in refused_retries) router.reset() except litellm.Timeout as e: pass From 7a7770db0d54e4734175fd7e798eb389ff92a1bd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 22:46:55 -0700 Subject: [PATCH 36/46] test(e2e): verify streamed answers and tool continuation --- tests/e2e/e2e_http.py | 2 + .../test_chat_stream_contract_e2e.py | 83 +++++++--- .../e2e/llm_translation/test_messages_e2e.py | 144 +++++++++++++++++- tests/e2e/models.py | 11 ++ .../test_streaming_iterator_tool_args.py | 8 +- 5 files changed, 218 insertions(+), 30 deletions(-) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 67370c98274..e0a20495964 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -170,6 +170,7 @@ class StreamingResponse(BaseModel): # the consumed body is elided, so this is the only place they surface. stream_error: str | None = None stream_done: bool = False + stream_done_positions: tuple[int, ...] = () @property def ok(self) -> bool: @@ -647,6 +648,7 @@ def streaming_outcome( stream_events=[payload for payload, _ in events], stream_event_arrivals=[arrived for _, arrived in events], stream_done=any(payload == _SSE_DONE for payload, _ in payloads), + stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE), stream_error=next( (line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)), None, diff --git a/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py b/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py index 4db2fe004c5..fdb76df703d 100644 --- a/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py +++ b/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py @@ -1,51 +1,90 @@ -"""Vendor §12.3: chat completions streaming SSE contract (LIT-4778). - -Asserts a streamed /chat/completions response is SSE, carries content chunks, -and terminates with the OpenAI [DONE] sentinel. -""" - from __future__ import annotations +from typing import Final + import pytest -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody, Usage from proxy_client import ProxyClient +from pydantic import BaseModel -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] + + +class _Delta(BaseModel): + content: str | None = None + + +class _Choice(BaseModel): + index: int + delta: _Delta + finish_reason: str | None = None + + +class _Chunk(BaseModel): + choices: tuple[_Choice, ...] + usage: Usage | None = None class TestChatStreamContract: @pytest.mark.covers("llm.chat_completions.openai.basic.stream.works") def test_chat_stream_is_sse_and_ends_with_done(self, proxy: ProxyClient, resources: ResourceManager) -> None: - model = f"e2e-chat-stream-{unique_marker()}" - model_id = proxy.create_model( + model: Final = f"e2e-chat-stream-{unique_marker()}" + base: Final = provider_edge_base("openai") + model_id: Final = proxy.create_model( model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), + LiteLLMParamsBody( + model="openai/gpt-5.6", + api_key="os.environ/OPENAI_API_KEY", + api_base=f"{base}/v1" if base else None, + ), ) resources.defer(lambda: proxy.delete_model(model_id)) - key = resources.key() - - result = proxy.chat_stream( + key: Final = resources.key() + expected: Final = "The amber kite crosses the quiet lake." + result: Final = proxy.chat_stream( key, ChatBody( model=model, messages=[ ChatMessage( - role="user", - content=f"Reply with the single word ok. {unique_marker()}", + role="user", content=f"Repeat exactly this sentence, with no additional text: {expected}" ) ], stream=True, - max_completion_tokens=32, - temperature=0.0, + stream_options=ChatStreamOptions(include_usage=True), + max_completion_tokens=256, + reasoning_effort="none", ), ) require_successful_call(result) assert result.is_streaming, f"expected SSE content-type, got {result.content_type!r}" assert result.stream_events, "stream returned no data events" - assert result.stream_done, ( - f"stream must terminate with [DONE]; " - f"chunks={result.chunks} done={result.stream_done} events={len(result.stream_events)}" + assert not result.stream_error, f"stream errored: {result.stream_error}" + assert result.stream_done, "stream must terminate with [DONE]" + assert result.stream_done_positions == (len(result.stream_events),), "[DONE] must occur once after all events" + chunks: Final = tuple(_Chunk.model_validate_json(event) for event in result.stream_events) + text_positions: Final = tuple( + i for i, chunk in enumerate(chunks) if any(c.delta.content for c in chunk.choices) ) + terminal_positions: Final = tuple( + i for i, chunk in enumerate(chunks) if any(c.finish_reason is not None for c in chunk.choices) + ) + assert text_positions, "stream completed without meaningful text" + assert len(terminal_positions) == 1, "expected exactly one terminal choice" + assert text_positions[0] < terminal_positions[0], "meaningful text must arrive before termination" + assert text_positions[-1] <= terminal_positions[0], "text arrived after termination" + assert all(c.index == 0 for chunk in chunks for c in chunk.choices) + assert tuple(c.finish_reason for c in chunks[terminal_positions[0]].choices) == ("stop",) + text: Final = "".join(c.delta.content or "" for chunk in chunks for c in chunk.choices) + assert text.strip() == expected, f"streamed answer was altered or incomplete: {text!r}" + usage_positions: Final = tuple(i for i, chunk in enumerate(chunks) if chunk.usage is not None) + assert usage_positions == (len(chunks) - 1,), "expected one final usage chunk" + assert terminal_positions[0] < usage_positions[0], "usage must follow the terminal choice" + usage: Final = chunks[-1].usage + assert usage is not None + assert usage.prompt_tokens is not None and usage.prompt_tokens > 0 + assert usage.completion_tokens is not None and usage.completion_tokens > 0 + assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index c731b52acd8..ca58c30d40c 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -21,7 +21,12 @@ from e2e_http import assert_client_error, require_successful_call, unwrap from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager from models import ( + AnthropicAssistantTurn, + AnthropicContentBlock, AnthropicCustomTool, + AnthropicToolChoice, + AnthropicToolResultBlock, + AnthropicToolResultTurn, AnthropicMessagesBody, ChatMessage, JsonSchemaProperty, @@ -29,7 +34,7 @@ from models import ( SpendLogRow, ToolInputSchema, ) -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @@ -284,8 +289,139 @@ class TestAnthropicMessages: result = endpoints_client.proxy.transport.send( "/v1/messages", headers=endpoints_client.proxy.transport.bearer(key), - json=_OptionalMessagesBody( - messages=[ChatMessage(role="user", content="hi")], max_tokens=50 - ), + json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50), ) assert_client_error(result, "messages missing model") + + +class _BridgeDelta(BaseModel): + type: str | None = None + partial_json: str | None = None + stop_reason: str | None = None + + +class _BridgeEvent(BaseModel): + type: str + index: int | None = None + content_block: AnthropicContentBlock | None = None + delta: _BridgeDelta | None = None + + +class _ParcelInput(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + parcel: str + shelf: int + + +def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock: + starts: Final = tuple( + event + for event in events + if event.type == "content_block_start" + and event.content_block is not None + and event.content_block.type == "tool_use" + ) + assert len(starts) == 1, "expected exactly one tool call" + start: Final = starts[0] + block: Final = start.content_block + assert block is not None and block.id and start.index is not None + fragments: Final = tuple( + event + for event in events + if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta" + ) + assert fragments, "tool stream contained no argument fragments" + assert all(event.index == start.index for event in fragments), "tool fragments changed index" + positions: Final = tuple(i for i, event in enumerate(events) if event in fragments) + stops: Final = tuple( + i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index + ) + assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0] + assert tuple( + event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None + ) == ("tool_use",) + terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta") + assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1 + assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), ( + "tool stream did not terminate exactly once" + ) + arguments: Final = _ParcelInput.model_validate_json( + "".join(event.delta.partial_json or "" for event in fragments if event.delta is not None) + ) + return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump()) + + +def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn: + assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call" + return AnthropicToolResultTurn(content=[result]) + + +def _request_tool( + client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool +) -> AnthropicContentBlock: + if stream: + response: Final = client.proxy.messages_stream(key, request) + require_successful_call(response) + assert response.is_streaming and not response.stream_error + return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events)) + response_body: Final = unwrap(client.proxy.messages(key, request)) + blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use") + assert len(blocks) == 1 + return blocks[0] + + +class TestOpenAIMessagesToolContinuation: + @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) + def test_required_tool_arguments_and_correlated_result( + self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool + ) -> None: + model: Final = f"e2e-bridge-tool-{unique_marker()}" + base: Final = provider_edge_base("openai") + model_id: Final = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key: Final = resources.key(models=[model]) + tool: Final = AnthropicCustomTool( + name="locate_parcel", + description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.", + input_schema=ToolInputSchema( + properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")}, + required=["parcel", "shelf"], + ), + ) + question: Final = ChatMessage( + role="user", + content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.", + ) + request: Final = AnthropicMessagesBody( + model=model, + max_tokens=2048, + messages=[question], + tools=[tool], + tool_choice=AnthropicToolChoice(type="tool", name=tool.name), + stream=stream, + ) + emitted: Final = _request_tool(endpoints_client, key, request, stream) + assert emitted.id and emitted.name == "locate_parcel" + assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed" + receipt: Final = f"receipt-{unique_marker()}" + result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt)) + continuation: Final = unwrap( + endpoints_client.proxy.messages( + key, + AnthropicMessagesBody( + model=model, + max_tokens=2048, + tools=[tool], + tool_choice=AnthropicToolChoice(type="none"), + messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn], + ), + ) + ) + answer: Final = "".join(block.text or "" for block in continuation.content or ()) + assert answer.strip() == receipt, "continuation did not consume the correlated tool result" + assert all(block.type != "tool_use" for block in continuation.content or ()) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 6fca1268ebc..ba6d0fe3334 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -283,10 +283,15 @@ class ChatToolResultTurn(BaseModel): type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn +class ChatStreamOptions(BaseModel): + include_usage: bool + + class ChatBody(BaseModel): model: str messages: Sequence[ChatTurn] stream: bool = False + stream_options: ChatStreamOptions | None = None max_tokens: int | None = None max_completion_tokens: int | None = None temperature: float | None = None @@ -488,12 +493,18 @@ class AnthropicToolResultTurn(BaseModel): type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn +class AnthropicToolChoice(BaseModel): + type: Literal["auto", "any", "tool", "none"] + name: str | None = None + + class AnthropicMessagesBody(BaseModel): model: str messages: list[AnthropicMessage] max_tokens: int stream: bool | None = None tools: list[AnthropicTool] | None = None + tool_choice: AnthropicToolChoice | None = None guardrails: list[str] | None = None cache: dict[str, bool] | None = {"no-cache": True} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py index 29e9279731d..a20aaf2e324 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py @@ -8,6 +8,8 @@ Without the fix, the AnthropicStreamWrapper silently dropped these arguments, causing tool_use blocks to arrive with empty input {}. """ +import json + from typing import List from unittest.mock import MagicMock @@ -139,9 +141,7 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args(): # Verify the delta carries the tool arguments delta_event = events[input_json_delta_idx] - assert delta_event["delta"][ - "partial_json" - ], "input_json_delta should have non-empty partial_json" + assert json.loads(delta_event["delta"]["partial_json"]) == {"location": "Boston"} @pytest.mark.asyncio @@ -300,7 +300,7 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args(): assert ( input_json_delta_idx == tool_start_idx + 1 ), "input_json_delta should immediately follow the tool_use content_block_start" - assert events[input_json_delta_idx]["delta"]["partial_json"] + assert json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"} def test_sync_stream_no_extra_delta_when_tool_args_empty(): From aaf924693a540572f90f512334ccbc056ce3554c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:13:50 -0700 Subject: [PATCH 37/46] fix(router): count num_retries_per_request across fallback hops num_retries_per_request has always capped the retries of one request with its fallback hops included. #40930 started reading the per-hop attempted_retries counter instead, and every fallback hop restarts that counter at zero, so a request could spend a fresh retry budget on each hop and the legacy fallback cap test started seeing the hop run. Router.log_retry now also keeps request_retry_count on the request metadata, incremented on every retry and fallback hop and never truncated the way previous_models is, and max_retries_per_request_hit reads that count. The flat retry records, the litellm_metadata coverage and caps above four from #40930 stay as they are, and the legacy test goes back to its previous_models == 0 assertion. --- litellm/__init__.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 4 +- litellm/router.py | 6 ++- tests/local_testing/test_router_fallbacks.py | 17 ++---- .../test_router_helper_utils.py | 10 ++-- .../rust_bridge/test_lifecycle.py | 11 ++-- tests/test_litellm/test_router.py | 52 +++++++++++++++++++ tests/test_litellm/test_utils.py | 13 ++--- 8 files changed, 85 insertions(+), 30 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 261457d6889..ccfbf80369f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop +num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### secret_manager_client: Optional[Any] = ( None # list of instantiated key management clients - e.g. azure kv, infisical, etc. diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 6e76bf9d49e..15380bc5d57 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) if not isinstance(metadata, Mapping): return False - attempted_retries: Final = metadata.get("attempted_retries") - return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries + retry_count: Final = metadata.get("request_retry_count") + return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count def get_or_create_metadata_bucket( diff --git a/litellm/router.py b/litellm/router.py index 1665583386f..ff50cac1328 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8378,7 +8378,8 @@ class Router: def log_retry(self, kwargs: dict, e: Exception) -> dict: """ - When a retry or fallback happens, record which model group, deployment and attempt just failed and why + When a retry or fallback happens, record which model group, deployment and attempt just failed and why, + and count it toward the request-wide num_retries_per_request cap """ from litellm.types.router import RetryAttemptRecord @@ -8402,7 +8403,10 @@ class Router: else () ) breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) + earlier_retry_count: Final = request_metadata.get("request_retry_count") + request_retry_count: Final = (earlier_retry_count if type(earlier_retry_count) is int else 0) + 1 kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict + kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 5d7955ad34a..82b832f89fd 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -5,7 +5,6 @@ import asyncio import os import time import traceback -from typing import Final import pytest @@ -14,7 +13,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.types.router import RetryAttemptRecord from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -23,15 +21,15 @@ class MyCustomHandler(CustomLogger): success: bool = False failure: bool = False previous_models: int = 0 - previous_model_records: tuple[RetryAttemptRecord, ...] = () def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call") print( f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}" ) - self.previous_model_records = tuple(kwargs["litellm_params"]["metadata"].get("previous_models", ())) - self.previous_models = len(self.previous_model_records) + self.previous_models = len( + kwargs["litellm_params"]["metadata"].get("previous_models", []) + ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} print(f"self.previous_models: {self.previous_models}") def log_post_api_call(self, kwargs, response_obj, start_time, end_time): @@ -720,14 +718,7 @@ async def test_async_fallbacks_max_retries_per_request(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - records: Final = customHandler.previous_model_records - assert customHandler.previous_models == len(records) - assert records - assert {record["model_group"] for record in records} == {"azure/gpt-3.5-turbo"} - assert next(record["exception_type"] for record in records if record["attempted_retries"] == 0) == "AuthenticationError" - refused_retries: Final = tuple(record for record in records if record["attempted_retries"]) - assert refused_retries - assert all("Max retries per request hit!" in record["exception_string"] for record in refused_retries) + assert customHandler.previous_models == 0 # 0 retries, 0 fallback router.reset() except litellm.Timeout as e: pass diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 5b06c5fdb01..c7e577366c3 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -631,9 +631,11 @@ def test_deployment_callback_respects_cooldown_time(model_list): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list, metadata_key): - """log_retry appends one flat record per failed attempt and copies neither the request kwargs nor - the request metadata into it""" + """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the + request metadata into it, and counts every failed attempt of the request independently of the + per-hop attempted_retries""" router = Router(model_list=model_list) + rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo") new_kwargs = router.log_retry( kwargs={ "model": "gpt-3.5-turbo", @@ -641,7 +643,7 @@ def test_log_retry(model_list, metadata_key): "messages": [{"role": "user", "content": "hi"}], metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"}, }, - e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"), + e=rate_limit_error, ) assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [ { @@ -652,6 +654,8 @@ def test_log_retry(model_list, metadata_key): "attempted_retries": 2, } ] + assert new_kwargs[metadata_key]["request_retry_count"] == 1 + assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2 def test_update_usage(model_list): diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/test_litellm/rust_bridge/test_lifecycle.py index 1f0b5591c2b..d73385621d5 100644 --- a/tests/test_litellm/rust_bridge/test_lifecycle.py +++ b/tests/test_litellm/rust_bridge/test_lifecycle.py @@ -8,7 +8,7 @@ from litellm.rust_bridge.lifecycle import check_limits @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize( - "cap, attempted_retries, refused", + "cap, request_retry_count, refused", [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], ids=[ "cap-above-four-reached", @@ -17,12 +17,15 @@ from litellm.rust_bridge.lifecycle import check_limits "cap-of-zero-refuses-first-retry", ], ) -def test_check_limits_reads_attempted_retries( - monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, attempted_retries: int, refused: bool +def test_check_limits_reads_request_retry_count( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool ) -> None: monkeypatch.setattr(litellm, "num_retries_per_request", cap) monkeypatch.setattr(litellm, "max_budget", None) - kwargs: Final = {"model": "mistral/mistral-ocr-latest", metadata_key: {"attempted_retries": attempted_retries}} + kwargs: Final = { + "model": "mistral/mistral-ocr-latest", + metadata_key: {"request_retry_count": request_retry_count}, + } if refused: with pytest.raises(RuntimeError, match="Max retries per request hit!"): check_limits(kwargs) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3cabcd71627..577332727d7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11066,6 +11066,58 @@ async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypa ] +def _failing_group_with_healthy_fallback_router(num_retries): + return litellm.Router( + model_list=[ + { + "model_name": "broken-group", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + }, + { + "model_name": "healthy-group", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake", "mock_response": "ok"}, + }, + ], + fallbacks=[{"broken-group": ["healthy-group"]}], + num_retries=num_retries, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cap, hop_refused", [(2, True), (4, False)], ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop"] +) +async def test_num_retries_per_request_counts_retries_across_fallback_hops(monkeypatch, cap, hop_refused): + """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a + fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero + and a request could spend far more retries than the cap allows.""" + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + router = _failing_group_with_healthy_fallback_router(num_retries=1) + recorder = _FallbackAttemptRecorder() + litellm.callbacks.append(recorder) + try: + request = router.acompletion(model="broken-group", messages=[{"role": "user", "content": "hi"}]) + if not hop_refused: + assert (await request).choices[0].message.content == "ok" + return + with pytest.raises(litellm.InternalServerError): + await request + finally: + litellm.callbacks.remove(recorder) + + assert recorder.failed_targets == ["healthy-group"] + hop_refusals = [ + record["attempted_retries"] + for record in recorder.breadcrumbs_per_target[0] + if record["model_group"] == "healthy-group" and "Max retries per request hit!" in record["exception_string"] + ] + assert hop_refusals == [0, 1] + + @pytest.mark.asyncio async def test_fallback_traceback_stays_available_at_debug_level(): """Dropping the stack from the ERROR line is only safe because the fallback path still diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d5feda6f892..3d60fe86d9c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4077,10 +4077,11 @@ class TestMetadataNoneHandling: _RETRY_CAP_CASES: Final = ( - pytest.param(5, {"attempted_retries": 5}, True, id="cap-above-four-reached"), - pytest.param(5, {"attempted_retries": 4}, False, id="cap-above-four-not-reached"), - pytest.param(0, {"attempted_retries": 0}, False, id="first-attempt-passes-cap-of-zero"), - pytest.param(0, {"attempted_retries": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(5, {"request_retry_count": 5}, True, id="cap-above-four-reached"), + pytest.param(5, {"request_retry_count": 4}, False, id="cap-above-four-not-reached"), + pytest.param(0, {"request_retry_count": 0}, False, id="first-attempt-passes-cap-of-zero"), + pytest.param(0, {"request_retry_count": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(0, {"attempted_retries": 1}, False, id="per-hop-attempted-retries-is-not-the-cap"), pytest.param(5, {"previous_models": ("a", "b", "c", "d", "e")}, False, id="breadcrumb-count-is-not-the-cap"), pytest.param(5, None, False, id="metadata-none"), ) @@ -4098,7 +4099,7 @@ def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): +def test_num_retries_per_request_reads_request_retry_count_sync(monkeypatch, metadata_key, cap, metadata, refused): monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: @@ -4111,7 +4112,7 @@ def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metad @pytest.mark.asyncio @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -async def test_num_retries_per_request_reads_attempted_retries_async(monkeypatch, metadata_key, cap, metadata, refused): +async def test_num_retries_per_request_reads_request_retry_count_async(monkeypatch, metadata_key, cap, metadata, refused): monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: From 2bf44ed35480b17d7757d8817985b73147a486f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:18:47 -0700 Subject: [PATCH 38/46] fix(guardrails): reject tool_use rewrites that are not JSON objects --- .../chat/guardrail_translation/handler.py | 36 ++++++++++++------- .../test_anthropic_guardrail_handler.py | 14 +++++--- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9f3c85b555a..b222548f4ec 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -234,19 +234,20 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge _TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) -def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: +def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None: + try: + return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments) + except ValidationError: + return None + + +def _write_back_tool_use( + message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object] +) -> None: content: Final = message.get("content", None) block: Final = content[target.content_idx] if isinstance(content, list) else None if not isinstance(block, dict): return - try: - rewritten_input: Final = _TOOL_USE_INPUT_ADAPTER.validate_json(shape.arguments) - except ValidationError: - verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned arguments that are not a JSON object for tool_use %s; keeping its input", - block.get("id"), - ) - return block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place if shape.name is not None and shape.name != block.get("name"): block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place @@ -688,6 +689,7 @@ class AnthropicMessagesHandler(BaseTranslation): scanned_tool_calls=scanned_tool_calls, pre_guardrail_tool_calls=pre_guardrail_tool_calls, returned_tool_calls=guardrailed_inputs.get("tool_calls"), + guardrail_name=guardrail_to_apply.guardrail_name, ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -1116,15 +1118,25 @@ class AnthropicMessagesHandler(BaseTranslation): scanned_tool_calls: tuple[ScannedToolCall, ...], pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], returned_tool_calls: Sequence[object] | None, + guardrail_name: str | None, ) -> None: post_guardrail_tool_calls: Final = _tool_call_shapes( returned_tool_calls if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) else tuple(item.tool_call for item in scanned_tool_calls) ) - for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): - if before != after: - _write_back_tool_use(messages[item.target.msg_idx], item.target, after) + rewritten: Final = tuple( + (item, after, _rewritten_tool_use_input(after.arguments)) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if before != after + ) + applicable: Final = tuple( + (item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None + ) + if len(applicable) != len(rewritten): + raise unappliable_request_rewrite(guardrail_name) + for item, after, rewritten_input in applicable: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input) async def process_output_response( self, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 287674afaa7..b73ef6453fa 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2328,16 +2328,20 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" @pytest.mark.asyncio - async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + async def test_non_json_rewritten_arguments_are_rejected_by_name(self): + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite + handler = AnthropicMessagesHandler() guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") data = self._tool_use_conversation(system="You are a careful agent harness.") + original = json.loads(json.dumps(data)) - await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert data["messages"][1]["content"][0]["input"] == { - "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" - } + assert excinfo.value.guardrail_name == "scan-only-capture" + assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" + assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" @pytest.mark.asyncio async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): From 4fca818f3483ca4d18a8efd7397a6f0892aa48c0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 23:45:48 -0700 Subject: [PATCH 39/46] fix(cli): align wide and combining Unicode cost labels --- .../client/cli/commands/statusline_script.py | 14 ++++++-- .../client/cli/test_statusline_script.py | 33 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 0e1c1b25e0f..815875d59b7 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -28,6 +28,7 @@ import os import sys import tempfile import time +import unicodedata import urllib.error import urllib.request from collections.abc import Callable, Mapping @@ -301,6 +302,14 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str: return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}" +def _display_width(label: str) -> int: + return sum( + 2 if unicodedata.east_asian_width(character) in ("W", "F") else 1 + for character in label + if unicodedata.category(character) not in ("Mn", "Me") + ) + + def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str: def paint(code: str, text: str) -> str: return f"{code}{text}{RESET}" if use_color else text @@ -315,13 +324,14 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") peak: Final = max(session.spend, session.baseline_spend) - label_width: Final = max(len(session.router_name), len(reference)) + label_width: Final = max(_display_width(session.router_name), _display_width(reference)) rows: Final = ( (session.router_name, session.spend, LITELLM_COLOR), (reference, session.baseline_spend, BASELINE_COLOR), ) lines: Final = ( - f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} " + f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} " + f"{_bar(amount / peak, color, bar_width, use_color)} " f"{paint(DIM, f'${amount:.2f}')}" for label, amount, color in rows ) diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 2b812932542..122c9601c33 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -265,6 +265,39 @@ class TestRender: "Claude Opus 5 ██████████ $0.38", ] + @pytest.mark.parametrize( + ("router_name", "baseline_name", "router_padding", "baseline_padding"), + ( + ("路由-router", "Claude Opus 5", 3, 1), + ("智能模型路由器", "Claude Opus 5", 1, 2), + ("ＡＢＣ-router", "Claude Opus 5", 1, 1), + ("cafe\u0301-router", "Claude Opus 5", 3, 1), + ("a\u20dd-router", "Claude Opus 5", 6, 1), + ("カ\u3099-router", "Claude Opus 5", 5, 1), + ("auto", "基準モデル", 7, 1), + ("auto", "cafe\u0301", 1, 1), + ), + ) + @pytest.mark.parametrize("use_color", (False, True)) + def test_unicode_labels_align_cost_bars_by_terminal_columns( + self, + config_dir: Path, + router_name: str, + baseline_name: str, + router_padding: int, + baseline_padding: int, + use_color: bool, + ) -> None: + (config_dir / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-opus-5", "display_name": baseline_name}]}) + ) + session: Final = RECORDED._replace(router_name=router_name) + text: Final = ANSI.sub("", render("claude-sonnet-5", session, config_dir, use_color, bar_width=10)) + assert text.splitlines()[1:] == [ + f"{router_name}{' ' * router_padding}████░░░░░░ $0.14", + f"{baseline_name}{' ' * baseline_padding}██████████ $0.38", + ] + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, From d0fdf1c2372756c4379362145a031bc4c1cf9fe4 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 23:54:50 -0700 Subject: [PATCH 40/46] fix(cli): show only the routed model in the footer header --- litellm/proxy/client/cli/README.md | 4 ++-- litellm/proxy/client/cli/commands/statusline_script.py | 8 ++------ .../proxy/client/cli/test_statusline_script.py | 10 +++++----- 3 files changed, 9 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index cb867cf9e61..a5a2675ed6e 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json` `lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline: ``` -claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 -LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14 +Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 +claude-auto ████████░░░░░░░░░░░░░░░░ $0.14 Claude Opus 5 ████████████████████████ $0.38 ``` diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 815875d59b7..47be3888a58 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -43,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3 BAR_WIDTH: Final = 24 BAR_FULL: Final = "\u2588" BAR_EMPTY: Final = "\u2591" -SEPARATOR: Final = " \u00b7 " TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024 CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",) CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY") @@ -315,11 +314,8 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo return f"{code}{text}{RESET}" if use_color else text routed: Final = paint(BOLD, f"Routed to: {model}") - if session is None: + if session is None or session.baseline_model is None or session.baseline_spend <= 0: return routed - header: Final = f"{session.router_name}{SEPARATOR}{routed}" - if session.baseline_model is None or session.baseline_spend <= 0: - return header reference: Final = baseline_label(session.baseline_model, config_dir) pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") @@ -335,7 +331,7 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo f"{paint(DIM, f'${amount:.2f}')}" for label, amount, color in rows ) - return "\n".join((f"{header} {delta}", *lines)) + return "\n".join((f"{routed} {delta}", *lines)) def color_enabled(env: Mapping[str, str]) -> bool: diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 122c9601c33..39d0e24d7b0 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -252,7 +252,7 @@ class TestRender: def test_savings_header_and_bars_against_the_routers_baseline(self, config_dir): text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) assert text.splitlines() == [ - "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", + "Routed to: claude-sonnet-5 -63% vs Claude Opus 5", "claude-auto ████░░░░░░ $0.14", "Claude Opus 5 ██████████ $0.38", ] @@ -330,7 +330,7 @@ class TestRender: assert "+25% vs Claude Opus 5" in render("m", dearer, config_dir, use_color=False) def test_without_a_baseline_only_the_routed_line_shows(self, config_dir): - assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "claude-auto · Routed to: m" + assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "Routed to: m" assert render("m", None, config_dir, False) == "Routed to: m" def test_color_wraps_the_same_text(self, config_dir): @@ -352,7 +352,7 @@ class TestClaudeCodeMode: return Fetched(RECORDED, definitive=True) text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) - assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert text.startswith("Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") assert text.splitlines()[1].startswith("claude-auto ") def test_a_discovered_display_name_labels_the_sessions_model( @@ -364,7 +364,7 @@ class TestClaudeCodeMode: return Fetched(session, definitive=True) text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) - assert text.startswith("claude-auto · Routed to: Claude Opus 5 -63% vs Claude Opus 5\n") + assert text.startswith("Routed to: Claude Opus 5 -63% vs Claude Opus 5\n") def test_an_unrecorded_session_degrades_to_the_routed_line(self, tmp_path, transcript, config_dir): assert _run(_payload(transcript), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( @@ -420,7 +420,7 @@ class TestCodexMode: out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) message = json.loads(out)["systemMessage"] - assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.splitlines()[1] == "Routed to: claude-sonnet-5 -63% vs Claude Opus 5" assert message.splitlines()[2].startswith("claude-auto ") assert message.startswith("\n") assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] From f80cb5cb4676a9d49009aef432cc9f5eda5ed15b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:04:02 -0700 Subject: [PATCH 41/46] fix(router): ignore planted request_retry_count seeds and cover the rust OCR cap path The router clamps a negative request_retry_count found in request metadata before counting a failure, and the proxy strips a client-supplied request_retry_count with the other router-reserved metadata fields. The rust OCR lifecycle test that trips the per-request cap now plants request_retry_count instead of attempted_retries, which the cap no longer reads since the previous commit --- litellm/proxy/litellm_pre_call_utils.py | 2 +- litellm/router.py | 4 ++-- .../test_router_helper_utils.py | 6 ++++-- .../proxy/test_litellm_pre_call_utils.py | 4 ++++ tests/test_litellm/test_router.py | 16 ++++++++++++---- tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- 6 files changed, 24 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 59971e54e46..14d0e7e478f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -334,7 +334,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( - {"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY} + {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY} ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" diff --git a/litellm/router.py b/litellm/router.py index ff50cac1328..cb0b7050876 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8403,8 +8403,8 @@ class Router: else () ) breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) - earlier_retry_count: Final = request_metadata.get("request_retry_count") - request_retry_count: Final = (earlier_retry_count if type(earlier_retry_count) is int else 0) + 1 + earlier: Final = request_metadata.get("request_retry_count") + request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index c7e577366c3..7949dd7818d 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -632,8 +632,8 @@ def test_deployment_callback_respects_cooldown_time(model_list): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list, metadata_key): """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the - request metadata into it, and counts every failed attempt of the request independently of the - per-hop attempted_retries""" + request metadata into it, counts every failed attempt of the request independently of the + per-hop attempted_retries, and never trusts a negative count planted before the first failure""" router = Router(model_list=model_list) rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo") new_kwargs = router.log_retry( @@ -656,6 +656,8 @@ def test_log_retry(model_list, metadata_key): ] assert new_kwargs[metadata_key]["request_retry_count"] == 1 assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2 + planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}} + assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1 def test_update_usage(model_list): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 37b983d709a..abd59bb4d3b 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7346,6 +7346,7 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: _PLANTED_STAMPS = { "attempted_fallbacks": 99, "original_model_group": "spoofed-group", + "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, "client_key": "client_value", } @@ -7378,6 +7379,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7403,6 +7405,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7431,6 +7434,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0} assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 577332727d7..1de86084980 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11089,18 +11089,26 @@ def _failing_group_with_healthy_fallback_router(num_retries): @pytest.mark.asyncio @pytest.mark.parametrize( - "cap, hop_refused", [(2, True), (4, False)], ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop"] + "cap, planted_count, hop_refused", + [(2, None, True), (4, None, False), (2, -100, True)], + ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop", "planted-negative-count-does-not-lift-the-cap"], ) -async def test_num_retries_per_request_counts_retries_across_fallback_hops(monkeypatch, cap, hop_refused): +async def test_num_retries_per_request_counts_retries_across_fallback_hops( + monkeypatch, cap, planted_count, hop_refused +): """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero - and a request could spend far more retries than the cap allows.""" + and a request could spend far more retries than the cap allows. A caller who plants a negative count + in the request metadata must not push the cap further away either.""" monkeypatch.setattr(litellm, "num_retries_per_request", cap) router = _failing_group_with_healthy_fallback_router(num_retries=1) recorder = _FallbackAttemptRecorder() litellm.callbacks.append(recorder) try: - request = router.acompletion(model="broken-group", messages=[{"role": "user", "content": "hi"}]) + metadata = {} if planted_count is None else {"request_retry_count": planted_count} + request = router.acompletion( + model="broken-group", messages=[{"role": "user", "content": "hi"}], metadata=metadata + ) if not hop_refused: assert (await request).choices[0].message.content == "ok" return diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 77d9ef167d0..dfcd63d3019 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -806,7 +806,7 @@ async def test_shared_call_limits_still_reject_before_reading_ocr_file( monkeypatch.setattr(litellm, "_current_cost", 2) monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError - arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"attempted_retries": 1}} + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"request_retry_count": 1}} with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) assert reads == [] From 92714cac0cffdf20ef612202605f19946f430f85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:07:45 -0700 Subject: [PATCH 42/46] fix(guardrails): validate tool_use rewrites before writing text rewrites back A guardrail that rewrites text and hands back tool_use arguments that are not a JSON object used to leave the text rewrite applied when the request was rejected, so failure logging saw a half-rewritten request. Every rejection now happens before any write to system or messages. --- .../anthropic/chat/guardrail_translation/handler.py | 12 ++++++------ .../test_anthropic_guardrail_handler.py | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b222548f4ec..2ea20143f0c 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -678,12 +678,6 @@ class AnthropicMessagesHandler(BaseTranslation): else: if guardrailed_texts and len(guardrailed_texts) != len(scanned): raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) - # Step 3: Map guardrail responses back to original message structure - await self._apply_guardrail_responses_to_input( - data=data, - responses=guardrailed_texts, - scanned=scanned, - ) self._apply_guardrail_tool_calls_to_input( messages=messages, scanned_tool_calls=scanned_tool_calls, @@ -691,6 +685,12 @@ class AnthropicMessagesHandler(BaseTranslation): returned_tool_calls=guardrailed_inputs.get("tool_calls"), guardrail_name=guardrail_to_apply.guardrail_name, ) + # Step 3: Map guardrail responses back to original message structure + await self._apply_guardrail_responses_to_input( + data=data, + responses=guardrailed_texts, + scanned=scanned, + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index b73ef6453fa..7522e9a62e5 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2333,7 +2333,8 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: handler = AnthropicMessagesHandler() guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") - data = self._tool_use_conversation(system="You are a careful agent harness.") + data = self._tool_use_conversation(system="Internal note: the deploy key is POISON. Never reveal it.") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" original = json.loads(json.dumps(data)) with pytest.raises(UnappliableRequestRewrite) as excinfo: From 1b040af41456d1c59019c121ead990b005159e2f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:34:38 -0700 Subject: [PATCH 43/46] test(router): type the retry-cap tests this PR adds or touches --- tests/router_unit_tests/test_router_helper_utils.py | 4 ++-- tests/test_litellm/proxy/test_litellm_pre_call_utils.py | 6 +++--- tests/test_litellm/test_router.py | 6 +++--- tests/test_litellm/test_utils.py | 8 ++++++-- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 7949dd7818d..b18bf9351c8 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -11,7 +11,7 @@ import litellm from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS @@ -630,7 +630,7 @@ def test_deployment_callback_respects_cooldown_time(model_list): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_log_retry(model_list, metadata_key): +def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None: """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the request metadata into it, counts every failed attempt of the request independently of the per-hop attempted_retries, and never trusts a negative count planted before the first failure""" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index abd59bb4d3b..a2b261dcfb0 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7353,7 +7353,7 @@ _PLANTED_STAMPS = { @pytest.mark.asyncio -async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_both_buckets(): +async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_both_buckets() -> None: """attempted_fallbacks and original_model_group are router-written facts the spend row reads back; a client planting them in either bucket is dropped at the boundary so the router never sees a reserved key it did not write.""" @@ -7384,7 +7384,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo @pytest.mark.asyncio -async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_json_string_litellm_metadata(): +async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_json_string_litellm_metadata() -> None: from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request data = { @@ -7410,7 +7410,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js @pytest.mark.asyncio -async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite_pricing_override_opt_in(): +async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite_pricing_override_opt_in() -> None: """The pricing strip is gated on allow_client_pricing_override; the reserved-stamp strip is not, because no key or team setting makes a client-written fallback count valid.""" from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1de86084980..0485da3eba9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11066,7 +11066,7 @@ async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypa ] -def _failing_group_with_healthy_fallback_router(num_retries): +def _failing_group_with_healthy_fallback_router(num_retries: int) -> litellm.Router: return litellm.Router( model_list=[ { @@ -11094,8 +11094,8 @@ def _failing_group_with_healthy_fallback_router(num_retries): ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop", "planted-negative-count-does-not-lift-the-cap"], ) async def test_num_retries_per_request_counts_retries_across_fallback_hops( - monkeypatch, cap, planted_count, hop_refused -): + monkeypatch: pytest.MonkeyPatch, cap: int, planted_count: int | None, hop_refused: bool +) -> None: """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero and a request could spend far more retries than the cap allows. A caller who plants a negative count diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3d60fe86d9c..e7878f6bff0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4099,7 +4099,9 @@ def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -def test_num_retries_per_request_reads_request_retry_count_sync(monkeypatch, metadata_key, cap, metadata, refused): +def test_num_retries_per_request_reads_request_retry_count_sync( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, metadata: object, refused: bool +) -> None: monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: @@ -4112,7 +4114,9 @@ def test_num_retries_per_request_reads_request_retry_count_sync(monkeypatch, met @pytest.mark.asyncio @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -async def test_num_retries_per_request_reads_request_retry_count_async(monkeypatch, metadata_key, cap, metadata, refused): +async def test_num_retries_per_request_reads_request_retry_count_async( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, metadata: object, refused: bool +) -> None: monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: From 0770f663c1c9300474c52d170aea81986c08e78b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:35:36 -0700 Subject: [PATCH 44/46] fix(vertex-live): report repeated search queries once per grounded turn The session usage collapsed duplicate query strings across turns while the price was per turn, so two turns asking the same question paid two fees yet reported web_search_requests 1. Sum each turn's grounding requests so the counter matches the bill; duplicates within one turn still collapse. --- ...tex_ai_live_passthrough_logging_handler.py | 36 +++++++++++-------- .../test_vertex_ai_live_passthrough.py | 30 ++++++++++++++++ 2 files changed, 52 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 57ec960b032..0ff02b29c58 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -13,6 +13,7 @@ from typing import Final, Literal, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.vertex_ai.gemini.grounding_requests import GroundingRequests, calculate_grounding_requests from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, ) @@ -28,6 +29,8 @@ from litellm.types.utils import ( Usage, ) +_NO_GROUNDING: Final = GroundingRequests(web_search_requests=None, google_maps_grounding_requests=None) + _AGGREGATED_FIELDS: Final = frozenset( { "promptTokenCount", @@ -72,6 +75,18 @@ def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], .. return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes))) +def _session_grounding_requests(websocket_messages: Sequence[object]) -> GroundingRequests: + per_turn: Final = tuple( + calculate_grounding_requests(_grounding_metadata(turn)) for turn in _turns(websocket_messages) + ) + web_search_requests: Final = sum(requests.web_search_requests or 0 for requests in per_turn) + google_maps_grounding_requests: Final = sum(requests.google_maps_grounding_requests or 0 for requests in per_turn) + return GroundingRequests( + web_search_requests=web_search_requests or None, + google_maps_grounding_requests=google_maps_grounding_requests or None, + ) + + _SummedField: TypeAlias = Literal[ "input_cost", "output_cost", @@ -223,7 +238,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): def _create_usage_object_from_metadata( usage_metadata: dict, model: str, - grounding_metadata: Sequence[Mapping[str, object]] = (), + grounding_requests: GroundingRequests = _NO_GROUNDING, ) -> Usage: """ Create a LiteLLM Usage object from Live API usage metadata. @@ -231,8 +246,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Args: usage_metadata: Usage metadata from the Live API response model: The model name - grounding_metadata: Every ``serverContent.groundingMetadata`` the session emitted, so - Search and Maps grounding carry their per-query charge + grounding_requests: The Search and Maps grounding requests summed over the session's + turns, matching the per-turn charge Returns: LiteLLM Usage object @@ -252,7 +267,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) - usage: Final = Usage( + return Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), @@ -262,6 +277,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): image_tokens=prompt_by_modality.get("IMAGE"), video_tokens=prompt_by_modality.get("VIDEO"), tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None, + web_search_requests=grounding_requests.web_search_requests, + google_maps_grounding_requests=grounding_requests.google_maps_grounding_requests, ), completion_tokens_details=CompletionTokensDetailsWrapper( text_tokens=candidates_by_modality.get("TEXT"), @@ -270,15 +287,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): video_tokens=candidates_by_modality.get("VIDEO"), ), ) - if grounding_metadata: - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet - usage, grounding_metadata - ) - return usage def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None: usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) @@ -286,7 +294,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): return None return self._create_usage_object_from_metadata( usage_metadata=usage_metadata, - grounding_metadata=_grounding_metadata(websocket_messages), + grounding_requests=_session_grounding_requests(websocket_messages), model=model, ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 70c2fda369b..8b3dc436b8f 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -534,6 +534,36 @@ class TestVertexAILivePassthroughLoggingHandler: assert two_breakdown["total_cost"] == pytest.approx(two_cost) assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_a_query_repeated_across_turns_is_reported_once_per_turn(self, handler): + """The reported query count must agree with the bill, which charges every grounded turn. + + The session usage collapsed duplicate query strings across turns while the price was + per turn, so two turns asking the same question paid two fees yet reported one query. + Duplicates within one turn still collapse, since that turn ran one search. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + logging_obj = self._priced_logging_obj() + + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=[head, grounding, turn, grounding, turn], + logging_obj=logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=self.NATIVE_AUDIO_MODEL, + custom_llm_provider="vertex_ai", + ) + _, one_breakdown = self._billed_session(handler, [head, grounding, turn]) + repeated_within_turn = handler._session_usage( + [head, self._grounding_frame({"webSearchQueries": ["q", "q"]}), turn], self.NATIVE_AUDIO_MODEL + ) + + assert result["result"].usage.prompt_tokens_details.web_search_requests == 2 + assert logging_obj.cost_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + assert repeated_within_turn.prompt_tokens_details.web_search_requests == 1 + def test_the_fixed_cost_margin_is_charged_once_per_session(self, handler): """A fixed cost margin is a flat per-request fee, and a Live session is one spend row. From 5969fbb052eb1cb12c3e6f66311832d8927da8bb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:44:46 -0700 Subject: [PATCH 45/46] refactor(realtime): freeze the empty credential mapping in the health check --- litellm/realtime_api/main.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 6fdadc52f13..f94ba4c4a02 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -60,7 +60,9 @@ _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: credential_name: Final = model_params.get("litellm_credential_name") credential_values: Final = ( - CredentialAccessor.get_credential_values(credential_name) if isinstance(credential_name, str) else {} + CredentialAccessor.get_credential_values(credential_name) + if isinstance(credential_name, str) + else _EMPTY_MODEL_PARAMS ) return MappingProxyType({**credential_values, **model_params}) From 87b630429d8684fdefb141fe08f80edc07269791 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:20:09 -0700 Subject: [PATCH 46/46] fix(realtime): send openai and xai health check keys as bearer tokens --- litellm/realtime_api/main.py | 17 ++++++---- tests/test_litellm/realtime_api/test_main.py | 34 ++++++++++++++++++++ 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f94ba4c4a02..d67e4555a29 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -55,6 +55,7 @@ xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_AUTH_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: @@ -602,13 +603,15 @@ def _azure_realtime_health_protocol( def _realtime_health_check_auth_headers( custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, Any] -) -> Mapping[str, str | None]: - if custom_llm_provider != "azure": - return MappingProxyType({"api-key": api_key}) - return azure_realtime.get_auth_headers( - api_key=api_key, - azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))), - ) +) -> Mapping[str, str]: + if custom_llm_provider == "azure": + return azure_realtime.get_auth_headers( + api_key=api_key, + azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))), + ) + if api_key is None: + return _EMPTY_AUTH_HEADERS + return MappingProxyType({"Authorization": f"Bearer {api_key}"}) async def _realtime_health_check( diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index dcd4c01730d..643e65af47c 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -276,6 +276,40 @@ async def test_azure_health_check_resolves_stored_credentials(monkeypatch): assert "api-version=2025-04-01-preview" in connect.url +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected_url"), + [ + ("xai", "grok-voice-latest", "wss://api.x.ai/v1/realtime?model=grok-voice-latest"), + ("openai", "gpt-realtime", "wss://api.openai.com/v1/realtime?model=gpt-realtime"), + ], +) +async def test_bearer_health_check_sends_stored_credential_as_bearer_token( + monkeypatch, custom_llm_provider: str, model: str, expected_url: str +): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="voice-key", + credential_values={"api_key": "sk-from-credential"}, + credential_info={}, + ) + ], + ) + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model=model, + custom_llm_provider=custom_llm_provider, + api_key=None, + model_params={"model": f"{custom_llm_provider}/{model}", "litellm_credential_name": "voice-key"}, + ) + assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer sk-from-credential"} + assert connect.url == expected_url + + @pytest.mark.asyncio async def test_azure_health_check_probes_ga_transcription_url_for_transcription_model(local_model_cost_map): """Regression for LIT-6240: transcription-only models (mode audio_transcription