From bfc7ae05bf78982b7cce01e5af37ad4b6e8574b5 Mon Sep 17 00:00:00 2001 From: unknown <> Date: Wed, 1 Jul 2026 14:20:26 +0000 Subject: [PATCH 001/157] feat: add claude-sonnet-5 pricing to model cost map Add Claude Sonnet 5 with introductory pricing (/0 per MTok, through Aug 31 2026) across Anthropic, Bedrock (global + regional), Vertex AI, Azure AI, and Snowflake providers. Includes cache pricing (1.25x write, 2x 1hr write, 0.1x read) and 128K output context. Fixes #31868 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + ...odel_prices_and_context_window_backup.json | 1836 +++++++++++------ model_prices_and_context_window.json | 1398 +++++++------ .../test_claude_sonnet_5_config.py | 163 ++ 4 files changed, 2137 insertions(+), 1261 deletions(-) create mode 100644 tests/test_litellm/test_claude_sonnet_5_config.py diff --git a/litellm/constants.py b/litellm/constants.py index aeb74a65839..1d17e7b6da6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1128,6 +1128,7 @@ BEDROCK_CONVERSE_MODELS = [ "anthropic.claude-opus-4-6-v1:0", "anthropic.claude-opus-4-6-v1", "anthropic.claude-sonnet-4-6", + "anthropic.claude-sonnet-5", "anthropic.claude-opus-4-1-20250805-v1:0", "anthropic.claude-opus-4-20250514-v1:0", "anthropic.claude-sonnet-4-20250514-v1:0", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 21132db93cb..15b16118a7b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1010,7 +1010,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1042,7 +1041,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1074,7 +1072,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1106,7 +1103,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1138,7 +1134,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1170,7 +1165,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1219,7 +1213,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1253,7 +1246,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1287,7 +1279,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1321,7 +1312,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1487,7 +1477,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1521,7 +1510,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1555,7 +1543,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1589,7 +1576,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1623,7 +1609,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1688,7 +1673,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1719,7 +1703,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1750,7 +1733,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1781,7 +1763,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1812,7 +1793,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1843,7 +1823,186 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "anthropic.claude-sonnet-5": { "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "global.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "us.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "eu.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "au.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "jp.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -2364,7 +2523,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2394,7 +2552,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2455,7 +2612,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2523,6 +2679,34 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, + "azure_ai/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, @@ -9132,17 +9316,16 @@ }, "bedrock/us-east-1/minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "output_cost_per_token": 1.2e-06 + "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -9754,17 +9937,16 @@ }, "bedrock/us-west-2/minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "output_cost_per_token": 1.2e-06 + "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -10274,6 +10456,35 @@ "supports_vision": true, "supports_output_config": true }, + "claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -16134,6 +16345,46 @@ "tpm": 8000000, "supports_image_size": false }, + "gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -16174,6 +16425,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -17588,6 +17877,48 @@ }, "supports_image_size": false }, + "gemini/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17630,6 +17961,47 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-image": { + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-flash-image-preview": { "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -18944,7 +19316,6 @@ "supported_endpoints": [ "/v1/chat/completions" ], - "supports_adaptive_thinking": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_vision": true @@ -21696,8 +22067,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21745,8 +22116,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21790,8 +22161,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21835,8 +22206,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21884,8 +22255,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21932,8 +22303,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21973,8 +22344,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22017,8 +22388,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22054,7 +22425,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -22062,8 +22433,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22100,7 +22471,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -22108,8 +22479,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22144,15 +22515,15 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22187,15 +22558,15 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22227,9 +22598,9 @@ "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, @@ -22263,9 +22634,9 @@ "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 272000, - "max_tokens": 272000, + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, @@ -24712,14 +25083,13 @@ }, "minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.2e-06, "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" @@ -28175,7 +28545,6 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_200k_tokens": 2.25e-05, "source": "https://openrouter.ai/anthropic/claude-sonnet-4.6", - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -28215,7 +28584,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -28277,7 +28645,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -29522,6 +29889,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/z-ai/glm-5.1": { + "input_cost_per_token": 1.05e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 5.25e-07, + "cache_creation_input_token_cost": 0.0, + "litellm_provider": "openrouter", + "max_input_tokens": 202752, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "source": "https://openrouter.ai/z-ai/glm-5.1", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/minimax/minimax-m2.1": { "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.2e-06, @@ -30200,7 +30583,6 @@ "supports_adaptive_thinking": true, "litellm_provider": "perplexity", "mode": "responses", - "supports_adaptive_thinking": true, "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, @@ -30210,7 +30592,6 @@ "supports_adaptive_thinking": true, "litellm_provider": "perplexity", "mode": "responses", - "supports_adaptive_thinking": true, "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, @@ -31126,7 +31507,7 @@ "supports_tool_choice": true }, "sambanova/Meta-Llama-3.2-1B-Instruct": { - "deprecation_date": "2025-06-25", + "deprecation_date": "2025-06-25", "input_cost_per_token": 4e-08, "litellm_provider": "sambanova", "max_input_tokens": 16384, @@ -31257,15 +31638,15 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "snowflake/claude-3-5-sonnet": { + "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "supports_computer_use": true, "supports_function_calling": true, "supports_vision": true, @@ -31273,14 +31654,14 @@ "supports_system_messages": true, "supports_response_schema": true }, - "snowflake/deepseek-r1": { + "snowflake/deepseek-r1": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000135, - "output_cost_per_token": 0.0000054, + "input_cost_per_token": 1.35e-06, + "output_cost_per_token": 5.4e-06, "supports_reasoning": true, "supports_system_messages": true }, @@ -31339,8 +31720,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.0000012, - "output_cost_per_token": 0.0000012, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, "supports_function_calling": true, "supports_system_messages": true }, @@ -31350,8 +31731,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "supports_function_calling": true, "supports_system_messages": true }, @@ -31361,8 +31742,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000024, - "output_cost_per_token": 0.00000024, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 2.4e-07, "supports_system_messages": true }, "snowflake/llama3.2-1b": { @@ -31379,17 +31760,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/llama3.3-70b": { + "snowflake/llama3.3-70b": { "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "snowflake", "mode": "chat", "supports_function_calling": true, "supports_system_messages": true - }, + }, "snowflake/mistral-7b": { "litellm_provider": "snowflake", "max_input_tokens": 32000, @@ -31404,14 +31785,14 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/mistral-large2": { + "snowflake/mistral-large2": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.000002, - "output_cost_per_token": 0.000006, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_response_schema": true @@ -31451,17 +31832,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/snowflake-llama-3.3-70b": { + "snowflake/snowflake-llama-3.3-70b": { "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "snowflake", "mode": "chat", "supports_function_calling": true, "supports_system_messages": true - }, + }, "stability/sd3": { "litellm_provider": "stability", "mode": "image_generation", @@ -32713,7 +33094,7 @@ "input_cost_per_token": 3.6e-06, "input_cost_per_token_above_200k_tokens": 7.2e-06, "output_cost_per_token_above_200k_tokens": 2.7e-05, - "cache_creation_input_token_cost_above_200k_tokens": 9.0e-06, + "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, "litellm_provider": "bedrock_converse", @@ -33455,7 +33836,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -34690,7 +35070,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34720,7 +35099,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34750,7 +35128,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34781,7 +35158,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34872,7 +35248,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34903,7 +35278,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34956,7 +35330,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -35265,6 +35638,21 @@ "tpm": 8000000, "supports_image_size": false }, + "vertex_ai/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -35280,6 +35668,19 @@ "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -37485,12 +37886,12 @@ }, "zai.glm-5": { "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.2e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -37511,20 +37912,6 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, - "zai.glm-5": { - "input_cost_per_token": 1e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 3.2e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "zai/glm-5": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2e-07, @@ -37540,6 +37927,21 @@ "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing" }, + "zai/glm-5.1": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-5-code": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 3e-07, @@ -37570,6 +37972,21 @@ "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing" }, + "zai/glm-4.7-flash": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 0, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 1.1e-07, @@ -42393,7 +42810,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -42411,6 +42827,64 @@ }, "supports_output_config": true }, + "vertex_ai/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, + "vertex_ai/claude-sonnet-5@default": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, "duckduckgo/search": { "litellm_provider": "duckduckgo", "mode": "search", @@ -42427,7 +42901,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -42442,7 +42919,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -42457,7 +42937,9 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -42471,7 +42953,9 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -42487,9 +42971,16 @@ "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/responses"], - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["text"], + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -42507,9 +42998,16 @@ "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/responses"], - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["text"], + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -42526,7 +43024,10 @@ "max_tokens": 256000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42542,7 +43043,10 @@ "max_tokens": 256000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42558,7 +43062,10 @@ "max_tokens": 128000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42717,20 +43224,6 @@ } ] }, - "zai.glm-5": { - "input_cost_per_token": 1e-06, - "output_cost_per_token": 3.2e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, @@ -42759,45 +43252,6 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, - "minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, - "bedrock/us-east-1/minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, - "bedrock/us-west-2/minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.5e-06, "cache_creation_input_token_cost_above_1hr": 2.4e-06, @@ -42844,362 +43298,382 @@ "supports_native_structured_output": true, "supports_pdf_input": true }, - "snowflake/claude-sonnet-4-5": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-sonnet-4-6": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-4-sonnet": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-4-opus": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000005, - "output_cost_per_token": 0.000025, - "cache_read_input_token_cost": 0.0000005, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/claude-haiku-4-5": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000005, - "cache_read_input_token_cost": 0.0000001, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-3-7-sonnet": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-4.1": { - "max_tokens": 16384, - "max_input_tokens": 300000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000002, - "output_cost_per_token": 0.000008, - "cache_read_input_token_cost": 0.0000005, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5": { - "max_tokens": 16384, - "max_input_tokens": 300000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000125, - "output_cost_per_token": 0.00001, - "cache_read_input_token_cost": 0.000000125, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5-mini": { - "max_tokens": 16384, - "max_input_tokens": 1000000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.0000003, - "output_cost_per_token": 0.0000012, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5-nano": { - "max_tokens": 16384, - "max_input_tokens": 5000000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000015, - "output_cost_per_token": 0.0000006, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/llama4-maverick": { - "max_tokens": 16384, - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000024, - "output_cost_per_token": 0.00000097, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true - }, - "snowflake/snowflake-arctic-embed-l-v2.0": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "input_cost_per_token": 0.00000007, - "output_cost_per_token": 0.0, - "litellm_provider": "snowflake", - "mode": "embedding" - }, - "snowflake/snowflake-arctic-embed-m-v2.0": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "input_cost_per_token": 0.00000007, - "output_cost_per_token": 0.0, - "litellm_provider": "snowflake", - "mode": "embedding" - }, - "soniox/stt-async-v4": { - "litellm_provider": "soniox", - "max_output_tokens": 8000, - "max_tokens": 8000, - "input_cost_per_second": 0.0, - "output_cost_per_second": 0.0000277778, - "mode": "audio_transcription", - "source": "https://soniox.com/pricing", - "supported_endpoints": ["/v1/audio/transcriptions"], - "supports_audio_input": true - }, - "soniox/stt-async-v5": { - "litellm_provider": "soniox", - "max_output_tokens": 8000, - "max_tokens": 8000, - "input_cost_per_second": 0.0, - "output_cost_per_second": 0.0000277778, - "mode": "audio_transcription", - "source": "https://soniox.com/pricing", - "supported_endpoints": ["/v1/audio/transcriptions"], - "supports_audio_input": true - }, - "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 6e-07, - "output_cost_per_token": 3.6e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 4.5e-07, - "output_cost_per_token": 1.8e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/Qwen/Qwen3.6-27B-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 3.2e-07, - "output_cost_per_token": 3.2e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 202752, - "max_output_tokens": 202752, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/moonshotai/Kimi-K2.6": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 9.6e-07, - "output_cost_per_token": 4e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/google/gemma-4-31B-it": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 5.6e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/openai/gpt-oss-120b": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/openai/gpt-oss-20b": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 7e-08, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - } - , + "snowflake/claude-sonnet-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-sonnet-4-6": { + "supports_adaptive_thinking": true, + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-sonnet-5": { + "supports_adaptive_thinking": true, + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-opus": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/claude-haiku-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "cache_read_input_token_cost": 1e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-3-7-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-4.1": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 8e-06, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.25e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-mini": { + "max_tokens": 16384, + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-nano": { + "max_tokens": 16384, + "max_input_tokens": 5000000, + "max_output_tokens": 16384, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/llama4-maverick": { + "max_tokens": 16384, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 9.7e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "snowflake/snowflake-arctic-embed-l-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + }, + "snowflake/snowflake-arctic-embed-m-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + }, + "soniox/stt-async-v4": { + "litellm_provider": "soniox", + "max_output_tokens": 8000, + "max_tokens": 8000, + "input_cost_per_second": 0.0, + "output_cost_per_second": 2.77778e-05, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supports_audio_input": true + }, + "soniox/stt-async-v5": { + "litellm_provider": "soniox", + "max_output_tokens": 8000, + "max_tokens": 8000, + "input_cost_per_second": 0.0, + "output_cost_per_second": 2.77778e-05, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 2.8e-09, @@ -43275,6 +43749,109 @@ "supports_tool_choice": true, "supports_vision": false }, + "deepseek/deepseek-v4-pro": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3.625e-09, + "input_cost_per_token": 4.35e-07, + "input_cost_per_token_cache_hit": 3.625e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "pinstripes/ps/glm-4.5-air": { + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "input_cost_per_token": 1.25e-07, + "output_cost_per_token": 4.5e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "source": "https://pinstripes.io/pricing" + }, + "pinstripes/ps/qwen3.6-35b-a3b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.5e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "source": "https://pinstripes.io/pricing" + }, + "pinstripes/ps/qwen3-30b-a3b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 2e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "source": "https://pinstripes.io/pricing" + }, + "pinstripes/ps/qwen3-coder-30b-a3b": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": false, + "source": "https://pinstripes.io/pricing" + }, + "pinstripes/ps/deepseek-v4-flash": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": true, + "source": "https://pinstripes.io/pricing" + }, + "pinstripes/ps/minimax-m2.7": { + "max_tokens": 1000192, + "max_input_tokens": 1000192, + "max_output_tokens": 1000192, + "input_cost_per_token": 2.55e-07, + "output_cost_per_token": 5.5e-07, + "litellm_provider": "pinstripes", + "mode": "chat", + "supports_function_calling": true, + "supports_assistant_prefill": true, + "supports_reasoning": false, + "source": "https://pinstripes.io/pricing" + }, "darkbloom/gemma-4-26b": { "input_cost_per_token": 3e-08, "litellm_provider": "darkbloom", @@ -43309,142 +43886,39 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "deepseek/deepseek-v4-pro": { - "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, - "litellm_provider": "deepseek", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://api-docs.deepseek.com/quick_start/pricing", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false - }, - "pinstripes/ps/glm-4.5-air": { - "max_tokens": 128000, - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "input_cost_per_token": 0.000000125, - "output_cost_per_token": 0.00000045, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" - }, - "pinstripes/ps/qwen3.6-35b-a3b": { - "max_tokens": 131072, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "input_cost_per_token": 0.00000014, - "output_cost_per_token": 0.00000045, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" - }, - "pinstripes/ps/qwen3-30b-a3b": { - "max_tokens": 131072, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "input_cost_per_token": 0.00000009, - "output_cost_per_token": 0.0000002, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" - }, - "pinstripes/ps/qwen3-coder-30b-a3b": { - "max_tokens": 131072, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "input_cost_per_token": 0.0000003, - "output_cost_per_token": 0.0000006, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" - }, - "pinstripes/ps/deepseek-v4-flash": { - "max_tokens": 163840, - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.0000002, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" - }, - "pinstripes/ps/minimax-m2.7": { - "max_tokens": 1000192, - "max_input_tokens": 1000192, - "max_output_tokens": 1000192, - "input_cost_per_token": 0.000000255, - "output_cost_per_token": 0.00000055, - "litellm_provider": "pinstripes", - "mode": "chat", - "supports_function_calling": true, - "supports_assistant_prefill": true, - "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" - }, - "fallback_generalizations": { - "rules": [ - { - "name": "anthropic-claude-adaptive-thinking", - "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", - "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", - "extends": "anthropic-claude", - "model_info": { - "supports_adaptive_thinking": true - } - }, - { - "name": "anthropic-claude", - "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", - "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", - "model_info": { - "litellm_provider": "anthropic", - "mode": "chat", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true, - "supports_tool_choice": true, - "supports_assistant_prefill": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_reasoning": true, - "supports_pdf_input": true, - "supports_system_messages": true - } - } - ] - } + "fallback_generalizations": { + "rules": [ + { + "name": "anthropic-claude-adaptive-thinking", + "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", + "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", + "extends": "anthropic-claude", + "model_info": { + "supports_adaptive_thinking": true + } + }, + { + "name": "anthropic-claude", + "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", + "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "model_info": { + "litellm_provider": "anthropic", + "mode": "chat", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_tool_choice": true, + "supports_assistant_prefill": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_system_messages": true + } + } + ] + } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 73cefeb7c77..15b16118a7b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1010,7 +1010,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1042,7 +1041,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1074,7 +1072,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1106,7 +1103,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1138,7 +1134,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1170,7 +1165,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1219,7 +1213,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1253,7 +1246,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1287,7 +1279,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1321,7 +1312,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1487,7 +1477,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1521,7 +1510,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1555,7 +1543,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1589,7 +1576,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1623,7 +1609,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1688,7 +1673,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1719,7 +1703,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1750,7 +1733,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1781,7 +1763,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1812,7 +1793,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -1843,7 +1823,186 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "anthropic.claude-sonnet-5": { "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "global.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "us.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "eu.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "au.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_output_config": true + }, + "jp.anthropic.claude-sonnet-5": { + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -2364,7 +2523,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2394,7 +2552,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2455,7 +2612,6 @@ "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -2523,6 +2679,34 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, + "azure_ai/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, @@ -9132,17 +9316,16 @@ }, "bedrock/us-east-1/minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "output_cost_per_token": 1.2e-06 + "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -9754,17 +9937,16 @@ }, "bedrock/us-west-2/minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "source": "https://aws.amazon.com/bedrock/pricing/", "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "output_cost_per_token": 1.2e-06 + "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -10274,6 +10456,35 @@ "supports_vision": true, "supports_output_config": true }, + "claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -19105,7 +19316,6 @@ "supported_endpoints": [ "/v1/chat/completions" ], - "supports_adaptive_thinking": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_vision": true @@ -21857,8 +22067,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21906,8 +22116,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21951,8 +22161,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21996,8 +22206,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22045,8 +22255,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22093,8 +22303,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22134,8 +22344,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22178,8 +22388,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22223,8 +22433,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22269,8 +22479,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22312,8 +22522,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22355,8 +22565,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -24873,14 +25083,13 @@ }, "minimax.minimax-m2.5": { "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.2e-06, "supports_function_calling": true, - "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" @@ -28336,7 +28545,6 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_200k_tokens": 2.25e-05, "source": "https://openrouter.ai/anthropic/claude-sonnet-4.6", - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -28376,7 +28584,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -28438,7 +28645,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -30377,7 +30583,6 @@ "supports_adaptive_thinking": true, "litellm_provider": "perplexity", "mode": "responses", - "supports_adaptive_thinking": true, "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, @@ -30387,7 +30592,6 @@ "supports_adaptive_thinking": true, "litellm_provider": "perplexity", "mode": "responses", - "supports_adaptive_thinking": true, "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, @@ -31434,15 +31638,15 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "snowflake/claude-3-5-sonnet": { + "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "supports_computer_use": true, "supports_function_calling": true, "supports_vision": true, @@ -31450,14 +31654,14 @@ "supports_system_messages": true, "supports_response_schema": true }, - "snowflake/deepseek-r1": { + "snowflake/deepseek-r1": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000135, - "output_cost_per_token": 0.0000054, + "input_cost_per_token": 1.35e-06, + "output_cost_per_token": 5.4e-06, "supports_reasoning": true, "supports_system_messages": true }, @@ -31516,8 +31720,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.0000012, - "output_cost_per_token": 0.0000012, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, "supports_function_calling": true, "supports_system_messages": true }, @@ -31527,8 +31731,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "supports_function_calling": true, "supports_system_messages": true }, @@ -31538,8 +31742,8 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.00000024, - "output_cost_per_token": 0.00000024, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 2.4e-07, "supports_system_messages": true }, "snowflake/llama3.2-1b": { @@ -31556,17 +31760,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/llama3.3-70b": { + "snowflake/llama3.3-70b": { "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "snowflake", "mode": "chat", "supports_function_calling": true, "supports_system_messages": true - }, + }, "snowflake/mistral-7b": { "litellm_provider": "snowflake", "max_input_tokens": 32000, @@ -31581,14 +31785,14 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/mistral-large2": { + "snowflake/mistral-large2": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "input_cost_per_token": 0.000002, - "output_cost_per_token": 0.000006, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_response_schema": true @@ -31628,17 +31832,17 @@ "max_tokens": 8192, "mode": "chat" }, - "snowflake/snowflake-llama-3.3-70b": { + "snowflake/snowflake-llama-3.3-70b": { "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, - "input_cost_per_token": 0.00000072, - "output_cost_per_token": 0.00000072, + "input_cost_per_token": 7.2e-07, + "output_cost_per_token": 7.2e-07, "litellm_provider": "snowflake", "mode": "chat", "supports_function_calling": true, "supports_system_messages": true - }, + }, "stability/sd3": { "litellm_provider": "stability", "mode": "image_generation", @@ -32890,7 +33094,7 @@ "input_cost_per_token": 3.6e-06, "input_cost_per_token_above_200k_tokens": 7.2e-06, "output_cost_per_token_above_200k_tokens": 2.7e-05, - "cache_creation_input_token_cost_above_200k_tokens": 9.0e-06, + "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, "litellm_provider": "bedrock_converse", @@ -33632,7 +33836,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -34867,7 +35070,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34897,7 +35099,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34927,7 +35128,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -34958,7 +35158,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -35049,7 +35248,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -35080,7 +35278,6 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -35133,7 +35330,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -37690,12 +37886,12 @@ }, "zai.glm-5": { "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.2e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -37716,20 +37912,6 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, - "zai.glm-5": { - "input_cost_per_token": 1e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 3.2e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "zai/glm-5": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 2e-07, @@ -42628,7 +42810,6 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -42646,6 +42827,64 @@ }, "supports_output_config": true }, + "vertex_ai/claude-sonnet-5": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, + "vertex_ai/claude-sonnet-5@default": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true + }, "duckduckgo/search": { "litellm_provider": "duckduckgo", "mode": "search", @@ -42662,7 +42901,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -42677,7 +42919,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, @@ -42692,7 +42937,9 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -42706,7 +42953,9 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -42722,9 +42971,16 @@ "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/responses"], - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["text"], + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -42742,9 +42998,16 @@ "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/responses"], - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["text"], + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -42761,7 +43024,10 @@ "max_tokens": 256000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42777,7 +43043,10 @@ "max_tokens": 256000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42793,7 +43062,10 @@ "max_tokens": 128000, "mode": "chat", "use_openai_responses_path": true, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_reasoning": true, @@ -42952,20 +43224,6 @@ } ] }, - "zai.glm-5": { - "input_cost_per_token": 1e-06, - "output_cost_per_token": 3.2e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, @@ -42994,45 +43252,6 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, - "minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, - "bedrock/us-east-1/minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, - "bedrock/us-west-2/minimax.minimax-m2.5": { - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "litellm_provider": "bedrock", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" - }, "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.5e-06, "cache_creation_input_token_cost_above_1hr": 2.4e-06, @@ -43079,363 +43298,382 @@ "supports_native_structured_output": true, "supports_pdf_input": true }, - "snowflake/claude-sonnet-4-5": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-sonnet-4-6": { - "supports_adaptive_thinking": true, - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-4-sonnet": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-4-opus": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000005, - "output_cost_per_token": 0.000025, - "cache_read_input_token_cost": 0.0000005, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/claude-haiku-4-5": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000005, - "cache_read_input_token_cost": 0.0000001, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/claude-3-7-sonnet": { - "max_tokens": 16384, - "max_input_tokens": 200000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, - "cache_read_input_token_cost": 0.0000003, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-4.1": { - "max_tokens": 16384, - "max_input_tokens": 300000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.000002, - "output_cost_per_token": 0.000008, - "cache_read_input_token_cost": 0.0000005, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5": { - "max_tokens": 16384, - "max_input_tokens": 300000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000125, - "output_cost_per_token": 0.00001, - "cache_read_input_token_cost": 0.000000125, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_vision": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5-mini": { - "max_tokens": 16384, - "max_input_tokens": 1000000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.0000003, - "output_cost_per_token": 0.0000012, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/openai-gpt-5-nano": { - "max_tokens": 16384, - "max_input_tokens": 5000000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000015, - "output_cost_per_token": 0.0000006, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true, - "supports_response_schema": true - }, - "snowflake/llama4-maverick": { - "max_tokens": 16384, - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "input_cost_per_token": 0.00000024, - "output_cost_per_token": 0.00000097, - "litellm_provider": "snowflake", - "mode": "chat", - "supports_function_calling": true, - "supports_system_messages": true - }, - "snowflake/snowflake-arctic-embed-l-v2.0": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "input_cost_per_token": 0.00000007, - "output_cost_per_token": 0.0, - "litellm_provider": "snowflake", - "mode": "embedding" - }, - "snowflake/snowflake-arctic-embed-m-v2.0": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "input_cost_per_token": 0.00000007, - "output_cost_per_token": 0.0, - "litellm_provider": "snowflake", - "mode": "embedding" - }, - "soniox/stt-async-v4": { - "litellm_provider": "soniox", - "max_output_tokens": 8000, - "max_tokens": 8000, - "input_cost_per_second": 0.0, - "output_cost_per_second": 0.0000277778, - "mode": "audio_transcription", - "source": "https://soniox.com/pricing", - "supported_endpoints": ["/v1/audio/transcriptions"], - "supports_audio_input": true - }, - "soniox/stt-async-v5": { - "litellm_provider": "soniox", - "max_output_tokens": 8000, - "max_tokens": 8000, - "input_cost_per_second": 0.0, - "output_cost_per_second": 0.0000277778, - "mode": "audio_transcription", - "source": "https://soniox.com/pricing", - "supported_endpoints": ["/v1/audio/transcriptions"], - "supports_audio_input": true - }, - "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 6e-07, - "output_cost_per_token": 3.6e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 4.5e-07, - "output_cost_per_token": 1.8e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/Qwen/Qwen3.6-27B-FP8": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 3.2e-07, - "output_cost_per_token": 3.2e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 202752, - "max_output_tokens": 202752, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/moonshotai/Kimi-K2.6": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 9.6e-07, - "output_cost_per_token": 4e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 0, - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/google/gemma-4-31B-it": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 5.6e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/openai/gpt-oss-120b": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - }, - "tensormesh/openai/gpt-oss-20b": { - "litellm_provider": "tensormesh", - "mode": "chat", - "input_cost_per_token": 7e-08, - "output_cost_per_token": 2.8e-07, - "cache_read_input_token_cost": 0, - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_reasoning": true, - "source": "https://serverless.tensormesh.ai/v1/models/openrouter" - } - , + "snowflake/claude-sonnet-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-sonnet-4-6": { + "supports_adaptive_thinking": true, + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-sonnet-5": { + "supports_adaptive_thinking": true, + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-4-opus": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/claude-haiku-4-5": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "cache_read_input_token_cost": 1e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/claude-3-7-sonnet": { + "max_tokens": 16384, + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-4.1": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 8e-06, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5": { + "max_tokens": 16384, + "max_input_tokens": 300000, + "max_output_tokens": 16384, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.25e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-mini": { + "max_tokens": 16384, + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/openai-gpt-5-nano": { + "max_tokens": 16384, + "max_input_tokens": 5000000, + "max_output_tokens": 16384, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "snowflake/llama4-maverick": { + "max_tokens": 16384, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "input_cost_per_token": 2.4e-07, + "output_cost_per_token": 9.7e-07, + "litellm_provider": "snowflake", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "snowflake/snowflake-arctic-embed-l-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + }, + "snowflake/snowflake-arctic-embed-m-v2.0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "snowflake", + "mode": "embedding" + }, + "soniox/stt-async-v4": { + "litellm_provider": "soniox", + "max_output_tokens": 8000, + "max_tokens": 8000, + "input_cost_per_second": 0.0, + "output_cost_per_second": 2.77778e-05, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supports_audio_input": true + }, + "soniox/stt-async-v5": { + "litellm_provider": "soniox", + "max_output_tokens": 8000, + "max_tokens": 8000, + "input_cost_per_second": 0.0, + "output_cost_per_second": 2.77778e-05, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 2.8e-09, @@ -43540,8 +43778,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.000000125, - "output_cost_per_token": 0.00000045, + "input_cost_per_token": 1.25e-07, + "output_cost_per_token": 4.5e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43553,8 +43791,8 @@ "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "input_cost_per_token": 0.00000014, - "output_cost_per_token": 0.00000045, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.5e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43566,8 +43804,8 @@ "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "input_cost_per_token": 0.00000009, - "output_cost_per_token": 0.0000002, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43579,8 +43817,8 @@ "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "input_cost_per_token": 0.0000003, - "output_cost_per_token": 0.0000006, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43592,8 +43830,8 @@ "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.0000002, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43605,8 +43843,8 @@ "max_tokens": 1000192, "max_input_tokens": 1000192, "max_output_tokens": 1000192, - "input_cost_per_token": 0.000000255, - "output_cost_per_token": 0.00000055, + "input_cost_per_token": 2.55e-07, + "output_cost_per_token": 5.5e-07, "litellm_provider": "pinstripes", "mode": "chat", "supports_function_calling": true, @@ -43648,39 +43886,39 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "fallback_generalizations": { - "rules": [ - { - "name": "anthropic-claude-adaptive-thinking", - "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", - "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", - "extends": "anthropic-claude", - "model_info": { - "supports_adaptive_thinking": true - } - }, - { - "name": "anthropic-claude", - "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", - "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", - "model_info": { - "litellm_provider": "anthropic", - "mode": "chat", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true, - "supports_tool_choice": true, - "supports_assistant_prefill": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_reasoning": true, - "supports_pdf_input": true, - "supports_system_messages": true - } - } - ] - } + "fallback_generalizations": { + "rules": [ + { + "name": "anthropic-claude-adaptive-thinking", + "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", + "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", + "extends": "anthropic-claude", + "model_info": { + "supports_adaptive_thinking": true + } + }, + { + "name": "anthropic-claude", + "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", + "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "model_info": { + "litellm_provider": "anthropic", + "mode": "chat", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_tool_choice": true, + "supports_assistant_prefill": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_system_messages": true + } + } + ] + } } diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py new file mode 100644 index 00000000000..943c4172205 --- /dev/null +++ b/tests/test_litellm/test_claude_sonnet_5_config.py @@ -0,0 +1,163 @@ +""" +Validate Claude Sonnet 5 model configuration entries. + +Sonnet 5 launched 2026-06-29 with introductory pricing ($2/$10 per MTok) +through August 31, 2026, after which it moves to standard pricing ($3/$15). +The cost-map entries use the introductory rates so spend tracking is accurate +during the promotional period. +""" + +import json +import os + +import pytest + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_sonnet_5_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = [ + ("claude-sonnet-5", "anthropic"), + ("anthropic.claude-sonnet-5", "bedrock_converse"), + ("vertex_ai/claude-sonnet-5", "vertex_ai-anthropic_models"), + ("azure_ai/claude-sonnet-5", "azure_ai"), + ] + + for model_name, provider in expected_models: + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + assert info["input_cost_per_token"] == 2e-06 + assert info["output_cost_per_token"] == 1e-05 + assert info["cache_creation_input_token_cost"] == 2.5e-06 + assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06 + assert info["cache_read_input_token_cost"] == 2e-07 + + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_adaptive_thinking"] is True + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + + +def test_sonnet_5_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + expected_models = { + "global.anthropic.claude-sonnet-5": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_creation_input_token_cost": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + }, + "us.anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + }, + "eu.anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + }, + "au.anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + }, + "jp.anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + for key, value in expected.items(): + assert info[key] == value, f"{model_name}.{key}: expected {value}, got {info[key]}" + + +def test_sonnet_5_present_in_bundled_backup(): + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in ( + "claude-sonnet-5", + "anthropic.claude-sonnet-5", + "global.anthropic.claude-sonnet-5", + "us.anthropic.claude-sonnet-5", + "eu.anthropic.claude-sonnet-5", + "au.anthropic.claude-sonnet-5", + "jp.anthropic.claude-sonnet-5", + "vertex_ai/claude-sonnet-5", + "vertex_ai/claude-sonnet-5@default", + "azure_ai/claude-sonnet-5", + "snowflake/claude-sonnet-5", + ): + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map): + info = litellm.get_model_info(model="claude-sonnet-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["input_cost_per_token"] == 2e-06 + assert info["output_cost_per_token"] == 1e-05 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_sonnet_5_all_variants_carry_adaptive_thinking_flag(cost_map): + variants = [k for k in cost_map if "claude-sonnet-5" in k and "sonnet-4-5" not in k] + assert variants, "no claude-sonnet-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" From 901a097b86ddb9539982173bcc91a6a29a27a2d7 Mon Sep 17 00:00:00 2001 From: unknown <> Date: Wed, 1 Jul 2026 14:35:36 +0000 Subject: [PATCH 002/157] fix: restore supports_reasoning on minimax entries lost during dedup The original JSON had duplicate minimax.minimax-m2.5 entries; the first had supports_reasoning: true, the second did not. JSON deduplication kept the second. Restore the flag on the surviving entries and add 1hr cache TTL assertions to the regional bedrock test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 9 ++++++--- model_prices_and_context_window.json | 9 ++++++--- tests/test_litellm/test_claude_sonnet_5_config.py | 5 +++++ 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 15b16118a7b..dba8bb949d0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9325,7 +9325,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -9946,7 +9947,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -25092,7 +25094,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "minimax/speech-02-hd": { "input_cost_per_character": 0.0001, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 15b16118a7b..dba8bb949d0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9325,7 +9325,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -9946,7 +9947,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -25092,7 +25094,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_reasoning": true }, "minimax/speech-02-hd": { "input_cost_per_character": 0.0001, diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py index 943c4172205..c8726a039f2 100644 --- a/tests/test_litellm/test_claude_sonnet_5_config.py +++ b/tests/test_litellm/test_claude_sonnet_5_config.py @@ -82,30 +82,35 @@ def test_sonnet_5_bedrock_regional_model_pricing(): "input_cost_per_token": 2e-06, "output_cost_per_token": 1e-05, "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, }, "us.anthropic.claude-sonnet-5": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 1.1e-05, "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, }, "eu.anthropic.claude-sonnet-5": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 1.1e-05, "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, }, "au.anthropic.claude-sonnet-5": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 1.1e-05, "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, }, "jp.anthropic.claude-sonnet-5": { "input_cost_per_token": 2.2e-06, "output_cost_per_token": 1.1e-05, "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, }, } From bbf51146ca7ffa63e5a46c903fa13b14b7f48c59 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 5 Sep 2026 13:29:38 +0000 Subject: [PATCH 003/157] fix(pricing): correct Vertex Haiku 4.5 output limit, Bedrock Mantle gpt-oss rates, add Scaleway deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 32 +++++++++++-------- model_prices_and_context_window.json | 32 +++++++++++-------- .../test_bedrock_mantle_transformation.py | 4 +-- 3 files changed, 40 insertions(+), 28 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2459ed940e0..1bc4c275161 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45421,8 +45421,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45446,8 +45446,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -51727,7 +51727,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -51766,7 +51767,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -51777,7 +51779,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -51800,7 +51803,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -51811,7 +51815,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -51833,7 +51838,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54283,7 +54289,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54317,8 +54323,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2459ed940e0..1bc4c275161 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45421,8 +45421,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45446,8 +45446,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -51727,7 +51727,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -51766,7 +51767,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -51777,7 +51779,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -51800,7 +51803,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -51811,7 +51815,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -51833,7 +51838,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54283,7 +54289,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54317,8 +54323,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index b88c27e64b9..92c012d339c 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -696,8 +696,8 @@ class TestBedrockMantlePricing: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") litellm.add_known_models() info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b") - # Bedrock pricing: $0.075/M input, $0.30/M output - assert info["input_cost_per_token"] == pytest.approx(7.5e-8) + # Bedrock pricing: $0.07/M input, $0.30/M output + assert info["input_cost_per_token"] == pytest.approx(7e-8) assert info["output_cost_per_token"] == pytest.approx(3e-7) def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch): From 528f3ae0161192004ca76d615363c748dcfdcd0d Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 5 Sep 2026 13:46:46 +0000 Subject: [PATCH 004/157] test(bedrock_mantle): cover gpt-oss-safeguard-20b input and output rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_mantle/test_bedrock_mantle_transformation.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 92c012d339c..b3eec97058a 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -700,6 +700,14 @@ class TestBedrockMantlePricing: assert info["input_cost_per_token"] == pytest.approx(7e-8) assert info["output_cost_per_token"] == pytest.approx(3e-7) + def test_gpt_oss_safeguard_20b_pricing(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") + litellm.add_known_models() + info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-safeguard-20b") + # Bedrock pricing: $0.07/M input, $0.20/M output + assert info["input_cost_per_token"] == pytest.approx(7e-8) + assert info["output_cost_per_token"] == pytest.approx(2e-7) + def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch): """ Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing. From b4e7776ab3534b4c6b91073aea1a1bbf8f681e83 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 5 Sep 2026 19:21:18 +0000 Subject: [PATCH 005/157] fix(model_prices): add gemini lyria-3.5, voyage-multilingual-2, chatgpt gpt-5.5 and gpt-5.6 entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 99 +++++++++++++++++++ model_prices_and_context_window.json | 99 +++++++++++++++++++ .../test_chatgpt_responses_transformation.py | 53 +++++++++- 3 files changed, 249 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1bc4c275161..8780c41b70a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27541,6 +27541,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -47654,6 +47718,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -60044,6 +60118,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_web_search": false, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1bc4c275161..8780c41b70a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27541,6 +27541,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -47654,6 +47718,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -60044,6 +60118,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_web_search": false, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8e0415d50de..8a0027fdb53 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -10,18 +10,23 @@ from unittest.mock import MagicMock, patch import httpx import pytest - +import litellm +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError +from litellm.main import responses_api_bridge_check from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager -from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig class TestChatGPTResponsesAPITransformation: @pytest.mark.parametrize( "model_name", [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", "chatgpt/gpt-5.4", "chatgpt/gpt-5.4-pro", "chatgpt/gpt-5.3-chat-latest", @@ -40,6 +45,50 @@ class TestChatGPTResponsesAPITransformation: assert isinstance(config, ChatGPTResponsesAPIConfig) assert config.custom_llm_provider == LlmProviders.CHATGPT + @pytest.mark.parametrize( + "model_name", + [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", + ], + ) + def test_chatgpt_responses_model_metadata(self, model_name, local_model_cost_map): + model_info = litellm.get_model_info(model_name) + + assert model_info["litellm_provider"] == "chatgpt" + assert model_info["mode"] == "responses" + assert model_info["supported_endpoints"] == [ + "/v1/chat/completions", + "/v1/responses", + ] + assert model_info["max_input_tokens"] == 1050000 + assert model_info["max_output_tokens"] == 128000 + + @pytest.mark.parametrize( + "model_name", + [ + "gpt-5.5", + "gpt-5.6-luna", + "gpt-5.6-sol", + "gpt-5.6-terra", + ], + ) + def test_chatgpt_models_bridge_chat_completions_to_responses(self, model_name, local_model_cost_map): + """A chat completions request for these models must take the Responses bridge. + + `gpt-5.6-*` also exists as an openai chat model, so an unregistered + chatgpt model resolves to mode "chat" here and never reaches the bridge. + """ + model_info, resolved_model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="chatgpt", + ) + + assert model_info["mode"] == "responses" + assert resolved_model == model_name + @patch("litellm.llms.chatgpt.responses.transformation.Authenticator") def test_chatgpt_responses_endpoint_url(self, mock_authenticator_class): mock_auth_instance = MagicMock() From 7015bf37bb3830dab79b3efd886eee9b5b7ffad6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 17:32:17 -0700 Subject: [PATCH 006/157] fix(proxy): apply team model aliases on the JWT auth path Team reads on the auth path never loaded the team's alias table, and the team-based JWT branch copied a hand-picked subset of team fields onto UserAPIKeyAuth, so aliases (and a few other team grants) never reached JWT callers: restricted teams 403'd alias requests and open teams 400'd them Load the alias relation where the team row is read and cached, project the team onto every team_* token field through one shared team_grants helper used by both JWT returns, and keep the relation when team model add/delete rewrites the cached team. ui_sso reuses the shared alias table model Resolves LIT-5858 Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- basedpyright-code-budget.json | 2 +- litellm/proxy/auth/auth_checks.py | 9 +- litellm/proxy/auth/handle_jwt.py | 5 +- litellm/proxy/auth/team_grants.py | 122 +++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 22 +-- .../management_endpoints/team_endpoints.py | 4 +- litellm/proxy/management_endpoints/ui_sso.py | 23 +--- .../proxy/auth/test_auth_checks.py | 64 +++++++++ .../proxy/auth/test_handle_jwt.py | 52 +++++++ .../proxy/auth/test_team_grants.py | 129 ++++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 116 ++++++++++++++++ .../test_team_endpoints.py | 44 ++++++ type-discipline-budget.json | 2 +- 13 files changed, 547 insertions(+), 47 deletions(-) create mode 100644 litellm/proxy/auth/team_grants.py create mode 100644 tests/test_litellm/proxy/auth/test_team_grants.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..8bdc251c684 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38269 }, "reportUnknownParameterType": { "limit": 19584 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc693317de0..a71a1993064 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -475,6 +475,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) +_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool: @@ -2858,7 +2859,9 @@ class TeamNotFoundError(HTTPException): async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) + response = await _team_table(TeamRepository(prisma_client)).find_unique( + where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + ) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -3158,7 +3161,9 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias}) + teams: Final = await _team_table(TeamRepository(prisma_client)).find_many( + where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + ) if not teams: raise HTTPException( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 0795cee7409..69091ee8344 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -53,6 +53,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -1595,7 +1596,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) ): is_allowed = allowed_routes_check( @@ -2132,7 +2133,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) except ProxyException: continue diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py new file mode 100644 index 00000000000..1196011dcdd --- /dev/null +++ b/litellm/proxy/auth/team_grants.py @@ -0,0 +1,122 @@ +"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``. + +The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path +starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT +callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through +``team_grants`` and the two paths cannot drift. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Annotated, Final + +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic.main import IncEx +from typing_extensions import ReadOnly, TypedDict + +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, +) + +_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) +_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType( + {"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})} +) + + +def _decode_model_aliases(value: object) -> object: + """``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both.""" + if not isinstance(value, str): + return value + try: + return _MODEL_ALIASES_ADAPTER.validate_json(value) + except ValidationError: + return None + + +class TeamModelAliasTable(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + + model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None + + +class _TeamJsonColumns(BaseModel): + """The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs.""" + + metadata: Mapping[str, object] | None = None + litellm_model_table: TeamModelAliasTable | None = None + + +class TeamGrants(TypedDict, total=False): + """Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply.""" + + team_alias: ReadOnly[str | None] + team_tpm_limit: ReadOnly[int | None] + team_rpm_limit: ReadOnly[int | None] + team_max_budget: ReadOnly[float | None] + team_soft_budget: ReadOnly[float | None] + team_spend: ReadOnly[float | None] + team_models: ReadOnly[Sequence[str]] + team_blocked: ReadOnly[bool] + team_metadata: ReadOnly[Mapping[str, object] | None] + team_model_aliases: ReadOnly[Mapping[str, str] | None] + team_object_permission_id: ReadOnly[str | None] + team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None] + team_member: ReadOnly[Member | None] + team_member_spend: ReadOnly[float | None] + team_member_tpm_limit: ReadOnly[int | None] + team_member_rpm_limit: ReadOnly[int | None] + + +def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: + try: + return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS)) + except ValidationError: + return _TeamJsonColumns() + + +def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None: + if team_object is None: + return None + alias_table: Final = _json_columns(team_object).litellm_model_table + return alias_table.model_aliases if alias_table is not None else None + + +def team_grants( + team_object: LiteLLM_TeamTable | None, + team_membership: LiteLLM_TeamMembership | None, + user_id: str | None, +) -> TeamGrants: + if team_object is None: + return TeamGrants() + json_columns: Final = _json_columns(team_object) + return TeamGrants( + team_alias=team_object.team_alias, + team_tpm_limit=team_object.tpm_limit, + team_rpm_limit=team_object.rpm_limit, + team_max_budget=team_object.max_budget, + team_soft_budget=team_object.soft_budget, + team_spend=team_object.spend, + team_models=tuple(team_object.models), + team_blocked=team_object.blocked, + team_metadata=json_columns.metadata, + team_model_aliases=( + json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None + ), + team_object_permission_id=team_object.object_permission_id, + team_object_permission=team_object.object_permission, + team_member=next( + (m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id), + None, + ), + team_member_spend=team_membership.spend if team_membership is not None else None, + team_member_tpm_limit=( + team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None + ), + team_member_rpm_limit=( + team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None + ), + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 93293db24c6..08dbe2508ec 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.resolvers import CredentialRef, Principal from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_grants from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( @@ -1476,24 +1477,16 @@ async def _user_api_key_auth_builder( user_id=user_id, user_email=user_email, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), - team_metadata=(team_object.metadata if team_object is not None else None), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, jwt_claims=jwt_claims, + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None and user_object.user_role is not None @@ -1507,17 +1500,8 @@ async def _user_api_key_auth_builder( user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), - team_member_rpm_limit=( - team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None - ), - team_member_tpm_limit=( - team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None - ), - team_metadata=(team_object.metadata if team_object is not None else None), jwt_claims=jwt_claims, - ) - valid_token.team_object_permission = ( - team_object.object_permission if team_object is not None else None + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ec68a10f65..c050368b3fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5601,7 +5601,7 @@ async def team_model_add( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( @@ -5688,7 +5688,7 @@ async def team_model_delete( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3e6434a5afd..c60888e298f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -22,7 +22,6 @@ from html import escape from types import MappingProxyType from typing import ( TYPE_CHECKING, - Annotated, Any, Final, Literal, @@ -42,7 +41,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -95,6 +94,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.team_grants import TeamModelAliasTable from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -209,31 +209,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]": return repo.table -_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) _SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -def _decode_model_aliases(value: object) -> object: - """``/team/new`` stores team model aliases as a JSON-encoded string in the Json column.""" - if not isinstance(value, str): - return value - try: - return _MODEL_ALIASES_ADAPTER.validate_json(value) - except ValidationError: - return None - - -class _TeamModelAliasTable(BaseModel): - model_config = ConfigDict(protected_namespaces=()) - - model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None - - class _TeamRowGrants(BaseModel): team_id: str team_alias: str | None = None models: tuple[str, ...] = () - litellm_model_table: _TeamModelAliasTable | None = None + litellm_model_table: TeamModelAliasTable | None = None class CliSsoTeamDetail(BaseModel): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2284a05b2e9..5bfef2b6445 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2374,6 +2374,44 @@ def _mock_prisma_for_team_lookup(find_unique): return mock_prisma_client +_TEAM_ALIAS_TABLE_ROW = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + +def _prisma_team_row(include): + """Mimics Prisma: the `litellm_model_table` relation rides on the row only when the query `include`s it.""" + columns = {"team_id": "team-aliases", "team_alias": "aliases", "models": ["gpt-4o"]} + row = ( + {**columns, "litellm_model_table": _TEAM_ALIAS_TABLE_ROW} + if (include or {}).get("litellm_model_table") + else columns + ) + return SimpleNamespace(dict=lambda: row, model_dump=lambda: row) + + +@pytest.mark.asyncio +async def test_get_team_object_loads_model_aliases_relation(): + """LIT-5858: the auth path read teams without `include`ing `litellm_model_table`, so every JWT + team came back with `model_aliases=None` and alias requests 403'd.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_unique(where, include=None): + return _prisma_team_row(include) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object( + team_id="team-aliases", + prisma_client=_mock_prisma_for_team_lookup(AsyncMock(side_effect=find_unique)), + user_api_key_cache=mock_cache, + check_db_only=True, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_team_object_distinguishes_absent_team_from_unreadable_row(): """A deleted team and a database that would not answer both surface as a 404, @@ -6195,6 +6233,32 @@ async def test_get_team_object_by_alias_db_fetch_returns_cached_obj(): assert result.models == ["gpt-4"] +@pytest.mark.asyncio +async def test_get_team_object_by_alias_loads_model_aliases_relation(): + """LIT-5858: same regression as `test_get_team_object_loads_model_aliases_relation`, for the + `team_alias_jwt_field` lookup.""" + from litellm.proxy.auth.auth_checks import get_team_object_by_alias + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_many(where, include=None): + return [_prisma_team_row(include)] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object_by_alias( + team_alias="aliases", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_org_object_by_alias_db_fetch_returns_validated_org(): from litellm.proxy._types import LiteLLM_OrganizationTable diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 99a0a4c0a8b..94226b5404d 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, LiteLLM_JWTAuth, + LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -1255,6 +1256,57 @@ async def test_find_team_with_model_access_model_group(monkeypatch): assert team_obj.team_id == "team-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_aliases", + ['{"fast": "gpt-4o"}', {"fast": "gpt-4o"}], + ids=["json-string", "dict"], +) +async def test_find_team_with_model_access_resolves_team_model_alias(monkeypatch, model_aliases): + """LIT-5858: a JWT team that grants `gpt-4o` under the alias `fast` must resolve a request + for `fast`. The JWT path used to pass `team_model_aliases=None`, so every alias request 403'd.""" + import sys + import types + + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + router = Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}]) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable( + team_id="team-aliases", + models=["gpt-4o"], + litellm_model_table=LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin"), + ) + + async def mock_get_team_object(*args, **kwargs): + return team + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + user_api_key_cache = DualCache() + + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids={"team-aliases"}, + requested_model="fast", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + ) + + assert team_id == "team-aliases" + assert team_obj is team + + @pytest.mark.asyncio async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatch): """Regression for #31189: a single-team JWT that grants the requested model diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py new file mode 100644 index 00000000000..447fc1c93a1 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -0,0 +1,129 @@ +import pytest + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_VerificationTokenView, + Member, + UserAPIKeyAuth, +) +from litellm.models.team import LiteLLM_ModelTable +from litellm.proxy.auth.team_grants import team_grants, team_model_aliases + +TEAM_ID = "team-grants" +USER_ID = "user-in-team" +ALIASES = {"fast": "gpt-4o-mini", "smart": "gpt-4o"} + + +def _alias_table(model_aliases) -> LiteLLM_ModelTable: + return LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin") + + +def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=TEAM_ID, + team_alias="grants-team", + tpm_limit=1000, + rpm_limit=10, + max_budget=50.0, + soft_budget=25.0, + spend=12.5, + models=["gpt-4o", "gpt-4o-mini"], + blocked=True, + metadata={"tier": "gold"}, + litellm_model_table=_alias_table(model_aliases), + object_permission_id="op-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["mcp-a"]), + members_with_roles=[ + Member(user_id="someone-else", role="user"), + Member(user_id=USER_ID, role="admin"), + ], + ) + + +def _membership() -> LiteLLM_TeamMembership: + return LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + spend=3.25, + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5), + ) + + +def test_team_grants_cover_every_team_field_the_key_path_gets(): + """Class guard for LIT-5858 and its siblings: every ``team_*`` column the combined-view SQL hands the + virtual-key path must come out of the projection too, with the team's actual value, so adding a column + to ``LiteLLM_VerificationTokenView`` without teaching ``team_grants`` fails here instead of in prod.""" + team = _full_team() + grants = team_grants(team_object=team, team_membership=_membership(), user_id=USER_ID) + token = UserAPIKeyAuth(team_id=TEAM_ID, **grants) + + view_team_fields = {name for name in LiteLLM_VerificationTokenView.model_fields if name.startswith("team_")} + assert view_team_fields - {"team_id"} <= set(grants) + assert all(grants[name] is not None for name in view_team_fields - {"team_id"}) + + assert token.team_alias == "grants-team" + assert token.team_tpm_limit == 1000 + assert token.team_rpm_limit == 10 + assert token.team_max_budget == 50.0 + assert token.team_soft_budget == 25.0 + assert token.team_spend == 12.5 + assert token.team_models == ["gpt-4o", "gpt-4o-mini"] + assert token.team_blocked is True + assert token.team_metadata == {"tier": "gold"} + assert token.team_model_aliases == ALIASES + assert token.team_object_permission_id == "op-1" + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_member == Member(user_id=USER_ID, role="admin") + assert token.team_member_spend == 3.25 + assert token.team_member_tpm_limit == 500 + assert token.team_member_rpm_limit == 5 + + +def test_team_grants_without_team_leave_token_defaults(): + token = UserAPIKeyAuth(**team_grants(team_object=None, team_membership=None, user_id=USER_ID)) + assert token == UserAPIKeyAuth() + + +@pytest.mark.parametrize( + "stored_aliases", + [ALIASES, '{"fast": "gpt-4o-mini", "smart": "gpt-4o"}'], + ids=["json-object", "json-string-as-written-by-team-new"], +) +def test_team_model_aliases_decode_both_storage_shapes(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) == ALIASES + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] == ALIASES + + +@pytest.mark.parametrize("stored_aliases", [None, "not json", '["a", "b"]', {"fast": 3}], ids=str) +def test_team_model_aliases_treat_unusable_column_as_no_aliases(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] is None + + +def test_team_model_aliases_none_without_relation_loaded(): + team = _full_team() + team.litellm_model_table = None + assert team_model_aliases(team) is None + assert team_model_aliases(None) is None + + +def test_team_member_is_the_callers_row_only(): + team = _full_team() + assert team_grants(team_object=team, team_membership=None, user_id="someone-else")["team_member"] == Member( + user_id="someone-else", role="user" + ) + assert team_grants(team_object=team, team_membership=None, user_id="stranger")["team_member"] is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_member"] is None + + +def test_membership_limits_absent_without_membership_row(): + grants = team_grants(team_object=_full_team(), team_membership=None, user_id=USER_ID) + assert grants["team_member_spend"] is None + assert grants["team_member_tpm_limit"] is None + assert grants["team_member_rpm_limit"] is None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index d44f96d95bf..78ffbb0db23 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -6872,3 +6872,119 @@ class TestLitellmReceivedAtStamping: assert result == earlier assert request.state.litellm_received_at == earlier + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"]) +async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin): + """LIT-5858: the team-based JWT path hand-built ``UserAPIKeyAuth`` from a short list of team fields, so the + team's model aliases (and on the admin return, its object permission) never reached the token and alias + requests 403'd. Both returns now go through ``team_grants``; pin the fields that used to be dropped.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.models.team import LiteLLM_ModelTable + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, + ) + + class _AcceptEveryJwt(JWTHandler): + def is_jwt(self, token: str) -> bool: + return True + + jwt_handler = _AcceptEveryJwt() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + team = LiteLLM_TeamTable( + team_id="team-jwt-aliases", + team_alias="jwt-aliases", + models=["gpt-4o"], + max_budget=40.0, + spend=4.0, + blocked=False, + metadata={"tier": "gold"}, + litellm_model_table=LiteLLM_ModelTable( + model_aliases='{"fast": "gpt-4o"}', created_by="admin", updated_by="admin" + ), + object_permission_id="op-jwt", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-jwt", mcp_servers=["mcp-a"]), + members_with_roles=[Member(user_id="jwt-user", role="admin")], + ) + membership = LiteLLM_TeamMembership(user_id="jwt-user", team_id="team-jwt-aliases", spend=1.5) + builder_result = { + "is_proxy_admin": is_proxy_admin, + "team_object": team, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": "jwt", + "team_id": "team-jwt-aliases", + "user_id": "jwt-user", + "user_email": "jwt-user@example.com", + "end_user_id": None, + "org_id": None, + "team_membership": membership, + "jwt_claims": {"sub": "jwt-user"}, + } + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {"enable_jwt_auth": True}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": jwt_handler, + "premium_user": True, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http", "headers": [], "method": "POST"}) + request._url = URL(url="/chat/completions") + with patch( # test-quality-ok: auth_builder is the claim-resolution seam; the regression is how its result is projected onto the token + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=builder_result, + ): + token = await _user_api_key_auth_builder( + request=request, + api_key="Bearer header.payload.signature", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert token.team_id == "team-jwt-aliases" + assert token.user_role == (LitellmUserRoles.PROXY_ADMIN if is_proxy_admin else LitellmUserRoles.INTERNAL_USER) + assert token.team_model_aliases == {"fast": "gpt-4o"} + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_object_permission_id == "op-jwt" + assert token.team_alias == "jwt-aliases" + assert token.team_models == ["gpt-4o"] + assert token.team_max_budget == 40.0 + assert token.team_spend == 4.0 + assert token.team_metadata == {"tier": "gold"} + assert token.team_member == Member(user_id="jwt-user", role="admin") + assert token.team_member_spend == 1.5 + assert token.jwt_claims == {"sub": "jwt-user"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 051e6bed4fd..2f6561046b1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2205,6 +2205,50 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): assert update_call_kwargs.get("include", {}).get("object_permission") is True +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["team_model_add", "team_model_delete"]) +async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_name, monkeypatch): + """LIT-5858: Prisma only returns `litellm_model_table` when the `update` asks for it, so the refreshed + cache entry lost the team's model aliases and JWT alias requests 403'd until the next DB read.""" + from litellm.proxy._types import TeamModelAddRequest, TeamModelDeleteRequest + from litellm.proxy.auth.team_grants import team_model_aliases + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_model_add, team_model_delete + + columns = {"team_id": "team-1234", "models": ["gpt-4o", "openai/*"]} + alias_table = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + async def update(where, data, include=None): + row = {**columns, "litellm_model_table": alias_table} if (include or {}).get("litellm_model_table") else columns + return SimpleNamespace(team_id="team-1234", model_dump=lambda: row) + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(model_dump=lambda: columns)) + prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=update) + prisma_client.db.execute_raw = AsyncMock(return_value=None) + cache = UserApiKeyCache() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + if endpoint_name == "team_model_add": + await team_model_add( + data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + else: + await team_model_delete( + data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + + cached_team = await cache.async_get_cache(key="team_id:team-1234", model_type=LiteLLM_TeamTableCachedObj) + assert team_model_aliases(cached_team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "endpoint_name", diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..7db8ed501f8 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22180 }, "LIT002": { - "limit": 26729 + "limit": 26727 }, "LIT003": { "limit": 261 From 8576cf74c2ef4eaa15f48eaaf99447f9d6fb394d Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Mon, 7 Sep 2026 14:09:37 +0200 Subject: [PATCH 007/157] feat(auth): add disable_env_credential_login setting with admin ui warning Env-credential login (UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset) is always live today. This adds a general_settings flag to turn that login path off once real admin accounts exist, and a warning banner shown to any admin while it remains enabled. The banner flag is served through /health/readiness/details and stays quiet when disable_password_login_when_sso_enabled already makes the env path unreachable. --- litellm/proxy/_types.py | 12 ++ litellm/proxy/auth/login_utils.py | 32 +++- .../health_endpoints/_health_endpoints.py | 20 +++ .../proxy/auth/test_login_utils.py | 154 ++++++++++++++++++ .../health_endpoints/test_health_endpoints.py | 27 +++ .../useHealthReadinessDetails.ts | 1 + .../src/app/(dashboard)/layout.test.tsx | 4 + .../src/app/(dashboard)/layout.tsx | 3 + .../EnvCredentialLoginWarningBanner.test.tsx | 76 +++++++++ .../EnvCredentialLoginWarningBanner.tsx | 35 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 11 files changed, 364 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c28ac8848ba..c30a3c5bea8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2817,6 +2817,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "UI username/password login. Default is False." ), ) + disable_env_credential_login: bool | None = Field( + None, + description=( + "If True, disables signing in to the Admin UI with the environment credentials: " + "UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback " + "means env-credential login is always live by default). Database users with passwords " + "are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password " + "before enabling, or nobody can sign in to the UI. A locked-out admin can still " + "administer the proxy over the API with the master key, and can unset this setting " + "and restart the proxy to restore env-credential login. Default is False." + ), + ) disable_budget_reservation: bool | None = Field( None, description=( diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 8d4f6f81363..6ca7ea5074d 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]: return ui_username, ui_password +def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool: + ui_username, ui_password = get_ui_credentials(master_key) + return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( + password.encode("utf-8"), ui_password.encode("utf-8") + ) + + +def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool: + """Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed. + + Two settings can turn it off: `disable_env_credential_login` unconditionally, and + `disable_password_login_when_sso_enabled` as a side effect, since its gate rejects + every username/password login before the env comparison runs. Feeds both the + `authenticate_user` gate and the Admin UI warning banner, so the banner never nags + about a login path that is already unreachable. + """ + if general_settings.get("disable_env_credential_login") is True: + return False + if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured(): + return False + return True + + class LoginResult: """Result object containing authentication data from login.""" @@ -129,7 +152,8 @@ async def authenticate_user( master_key: Master key for the proxy (required) prisma_client: Prisma database client (optional) general_settings: Proxy general_settings, checked for - `disable_password_login_when_sso_enabled` + `disable_password_login_when_sso_enabled` and + `disable_env_credential_login` Returns: LoginResult: Object containing authentication data @@ -170,8 +194,6 @@ async def authenticate_user( code=500, ) - ui_username, ui_password = get_ui_credentials(master_key) - # Check if we can find the `username` in the db. On the UI, users can enter username=their email _user_row: LiteLLM_UserTable | None = None user_role: ( @@ -197,8 +219,8 @@ async def authenticate_user( - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( - password.encode("utf-8"), ui_password.encode("utf-8") + if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials( + username, password, master_key ): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1785a2f0992..e5e4f89234b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1582,6 +1582,23 @@ async def _show_no_redis_warning() -> bool: return await count_live_proxy_workers(prisma_client) != 1 +def _show_env_credential_login_warning() -> bool: + """ + Whether the UI should warn admins that env-credential login is still enabled. + + UI_USERNAME/UI_PASSWORD (or the master key, when UI_PASSWORD is unset) grant + proxy-admin access with a shared static secret: no per-person identity, no + audit trail, no password policy, and it stays valid until the env var or + master key rotates. That is fine for first-time setup, so it is on by + default, but once real admin accounts exist it should be turned off with + `general_settings.disable_env_credential_login`. + """ + from litellm.proxy.auth.login_utils import is_env_credential_login_enabled + from litellm.proxy.proxy_server import general_settings + + return is_env_credential_login_enabled(general_settings) + + async def _get_health_readiness_details( response: Response | None = None, ) -> dict[str, Any]: @@ -1623,6 +1640,7 @@ async def _get_health_readiness_details( log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) show_no_redis_warning: Final = await _show_no_redis_warning() + show_env_credential_login_warning: Final = _show_env_credential_login_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected @@ -1650,6 +1668,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } else: return { @@ -1662,6 +1681,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } except Exception as e: raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})") diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 8d93d801bfd..72cc6e04a12 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.login_utils import ( LoginResult, authenticate_user, get_ui_credentials, + is_env_credential_login_enabled, ) @@ -799,3 +800,156 @@ class TestDisablePasswordLoginWhenSSOEnabled: assert isinstance(result, LoginResult) assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestDisableEnvCredentialLogin: + """`disable_env_credential_login` must reject a login with the env + credentials (UI_USERNAME/UI_PASSWORD, or the master-key fallback when + UI_PASSWORD is unset) while leaving database-user password logins + untouched, so admins with real accounts keep a way in.""" + + @pytest.mark.asyncio + async def test_rejects_correct_env_credentials_when_disabled(self): + master_key = "sk-1234" + ui_username = "admin" + ui_password = "env-only-password" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_rejects_master_key_fallback_when_disabled(self): + """With UI_PASSWORD unset, the master key IS the env password, so the + setting must reject it too or it protects nothing by default.""" + master_key = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": "admin"}, clear=True): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username="admin", + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_db_user_login_still_works_when_disabled(self): + master_key = "sk-1234" + user_email = "admin@example.com" + password = "Str0ng!Passw0rd" + + mock_user = LiteLLM_UserTable( + user_id="db-admin-1", + user_email=user_email, + password=hash_token(token=password), + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) + + with patch.dict( + os.environ, + { + "UI_USERNAME": "admin", + "UI_PASSWORD": "env-password", + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "db-user-token"}, + ) + ) + result = await authenticate_user( + username=user_email, + password=password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "db-admin-1" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + @pytest.mark.asyncio + async def test_env_login_still_works_when_setting_absent(self): + """Env-credential login is the bootstrap path on a fresh install and + must stay on by default.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestIsEnvCredentialLoginEnabled: + """Drives the Admin UI warning banner: it must be True exactly when a + login with the env credentials could actually succeed.""" + + def test_enabled_by_default(self): + assert is_env_credential_login_enabled({}) is True + + def test_disabled_by_dedicated_setting(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": True}) is False + + def test_explicit_false_keeps_it_enabled(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": False}) is True + + def test_disabled_when_sso_gate_blocks_all_password_logins(self): + """`disable_password_login_when_sso_enabled` with SSO configured + rejects every username/password login before the env comparison runs, + so the banner must not nag about an already-unreachable path.""" + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is False + + def test_enabled_when_sso_gate_is_set_but_sso_not_configured(self): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=False) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e9e58347337..b06d87ac67e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1301,6 +1301,33 @@ def test_health_readiness_details_returns_diagnostic_fields(monkeypatch): assert "cache" in response_data +@pytest.mark.parametrize( + "general_settings, expected_warning", + [ + ({}, True), + ({"disable_env_credential_login": True}, False), + ], +) +def test_health_readiness_details_reports_env_credential_login_warning(monkeypatch, general_settings, expected_warning): + """ + The Admin UI banner is driven by this flag: it must be True while + env-credential login is possible and False once + `disable_env_credential_login` turns that login path off. + """ + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/health/readiness/details") + + assert response.status_code == 200, response.text + assert response.json()["show_env_credential_login_warning"] is expected_warning + + def test_health_readiness_allows_explicit_legacy_public_details(monkeypatch): """ Operators can explicitly preserve the legacy public readiness payload. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts index 307fa9e1691..44d9092df34 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts @@ -14,6 +14,7 @@ export interface HealthReadinessDetailsResponse { log_level?: string; is_detailed_debug?: boolean; show_no_redis_warning?: boolean; + show_env_credential_login_warning?: boolean; } const fetchHealthReadinessDetails = async (accessToken: string): Promise => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7f1f4cc4bd5..3fe34610260 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -29,6 +29,10 @@ vi.mock("@/components/NoRedisWarningBanner", () => ({ NoRedisWarningBanner: () => null, })); +vi.mock("@/components/EnvCredentialLoginWarningBanner", () => ({ + EnvCredentialLoginWarningBanner: () => null, +})); + vi.mock("@/components/LicenseExpiryBanner", () => ({ LicenseExpiryBanner: () => null, })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 98f2a36d6f3..fa6df7f176a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -10,6 +10,7 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { useRouter, useSearchParams } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; +import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; import { uiHref } from "@/utils/uiHref"; @@ -113,6 +114,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -132,6 +134,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx new file mode 100644 index 00000000000..535694f14da --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx @@ -0,0 +1,76 @@ +import { renderWithProviders, screen } from "../../tests/test-utils"; +import { vi } from "vitest"; +import { EnvCredentialLoginWarningBanner } from "./EnvCredentialLoginWarningBanner"; +import type { HealthReadinessDetailsResponse } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import type { UseQueryResult } from "@tanstack/react-query"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ + useHealthReadinessDetails: vi.fn(), +})); +vi.mock("@/contexts/AuthContext", () => ({ + useAuth: vi.fn(), +})); + +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; + +const mockDetails = (data: Partial | undefined) => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data } as UseQueryResult); +}; + +const mockRole = (userRole: string) => { + vi.mocked(useAuth).mockReturnValue({ userRole } as ReturnType); +}; + +describe("EnvCredentialLoginWarningBanner", () => { + it("should warn an admin when env-credential login is enabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("Environment-credential login is enabled")).toBeInTheDocument(); + }); + + it("should tell the admin to create a regular admin account before disabling", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByText(/First create a regular admin account/i)).toBeInTheDocument(); + expect(screen.getByText("general_settings.disable_env_credential_login: true")).toBeInTheDocument(); + }); + + it("should warn an admin viewer too", () => { + mockRole("Admin Viewer"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + }); + + it("should render nothing for a non-admin even when the proxy reports the warning", () => { + mockRole("Internal User"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when env-credential login is disabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: false }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when readiness details are unavailable", () => { + mockRole("Admin"); + mockDetails(undefined); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should pass the access token to the readiness hook", () => { + mockRole("Admin"); + mockDetails(undefined); + renderWithProviders(); + expect(useHealthReadinessDetails).toHaveBeenCalledWith("my-token"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx new file mode 100644 index 00000000000..3a9d50011f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx @@ -0,0 +1,35 @@ +"use client"; + +import React from "react"; +import { TriangleAlert } from "lucide-react"; +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; +import { isAdminRole } from "@/utils/roles"; + +export const EnvCredentialLoginWarningBanner: React.FC<{ accessToken: string | null }> = ({ accessToken }) => { + const { userRole } = useAuth(); + const { data: healthData } = useHealthReadinessDetails(accessToken); + + if (!isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) { + return null; + } + + return ( +
+
+ ); +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..dae0af3d3b1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25778,6 +25778,11 @@ export interface components { * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed. */ disable_budget_reservation?: boolean | null; + /** + * Disable Env Credential Login + * @description If True, disables signing in to the Admin UI with the environment credentials: UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback means env-credential login is always live by default). Database users with passwords are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password before enabling, or nobody can sign in to the UI. A locked-out admin can still administer the proxy over the API with the master key, and can unset this setting and restart the proxy to restore env-credential login. Default is False. + */ + disable_env_credential_login?: boolean | null; /** * Disable Password Login When Sso Enabled * @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False. From 15c9fca53fc04c5ab05ebba0a8f047ef77394abd Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 7 Sep 2026 13:17:30 +0000 Subject: [PATCH 008/157] fix(pricing): add xai grok-imagine-video entries, computer-use-preview deprecation, gemini live 2.5 native audio card values Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 94 +++++++++++++++++-- model_prices_and_context_window.json | 94 +++++++++++++++++-- 2 files changed, 174 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c558a953032..33e9ce94455 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14336,6 +14336,7 @@ "supports_tool_choice": true }, "computer-use-preview": { + "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, "litellm_provider": "azure", "max_input_tokens": 8192, @@ -23829,9 +23830,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -23854,12 +23855,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -59479,6 +59480,85 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c558a953032..33e9ce94455 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14336,6 +14336,7 @@ "supports_tool_choice": true }, "computer-use-preview": { + "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, "litellm_provider": "azure", "max_input_tokens": 8192, @@ -23829,9 +23830,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -23854,12 +23855,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -59479,6 +59480,85 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", From 68d7a2d98fadd8f15eec9d5b0005e79de33e98bb Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 7 Sep 2026 13:29:29 +0000 Subject: [PATCH 009/157] fix(pricing): drop supported_endpoints from xai grok-imagine-video entries to match video model convention Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ------------ model_prices_and_context_window.json | 12 ------------ 2 files changed, 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 33e9ce94455..537db678f60 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59487,9 +59487,6 @@ "output_cost_per_second": 0.05, "output_cost_per_second_480p": 0.05, "source": "https://docs.x.ai/docs/models/grok-imagine-video", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59507,9 +59504,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59527,9 +59521,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59547,9 +59538,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 33e9ce94455..537db678f60 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59487,9 +59487,6 @@ "output_cost_per_second": 0.05, "output_cost_per_second_480p": 0.05, "source": "https://docs.x.ai/docs/models/grok-imagine-video", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59507,9 +59504,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59527,9 +59521,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", @@ -59547,9 +59538,6 @@ "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", - "supported_endpoints": [ - "/v1/videos" - ], "supported_modalities": [ "text", "image", From 6c21be619441dbd24879e1f8a4b897c6dbd667fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:15:06 -0700 Subject: [PATCH 010/157] fix(e2e): clean up batch files reliably and expire Azure inputs --- tests/e2e/batches/COVERAGE.md | 18 ++ tests/e2e/batches/batch_cleanup.py | 90 +++++++++ tests/e2e/batches/batch_client.py | 16 +- tests/e2e/batches/capabilities.py | 4 + tests/e2e/batches/conftest.py | 10 +- tests/e2e/batches/test_batch_cleanup.py | 187 ++++++++++++++++++ tests/e2e/batches/test_batches_e2e.py | 80 ++++---- .../test_managed_files_enforcement_e2e.py | 3 +- tests/e2e/lifecycle.py | 23 ++- 9 files changed, 381 insertions(+), 50 deletions(-) create mode 100644 tests/e2e/batches/batch_cleanup.py create mode 100644 tests/e2e/batches/test_batch_cleanup.py diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 8a7b68511ec..899530dde2b 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -120,6 +120,24 @@ create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: nested managed ids round-trip retrieve. This self-chaining only needs the proxy to reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. +## Cleanup + +Batch teardown cancels active batches before deleting their input files and keys. +Raw file IDs from both `model_param` and `provider_fallback` uploads use the upload +provider when deleted. Model-encoded and managed file IDs route themselves + +File deletion and batch cancellation check their responses and retry transient +failures up to three times. Teardown attempts every registered cleanup before +reporting failures as test errors. Already deleted files and batches that are +terminal are safe to clean up again. Cancellation polls for up to ten minutes +before input deletion, because accepting cancellation does not finish it + +Azure input uploads request `expires_after` anchored to `created_at` with +`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a +fallback for interrupted runs: immediate deletion remains the normal cleanup. +Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot +be requested through its Files API + ## Terminal state + cost write-back (cross-run marker baton) The 24h completion window rules out submit-and-wait inside one run, so diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py new file mode 100644 index 00000000000..a5b5e9bba37 --- /dev/null +++ b/tests/e2e/batches/batch_cleanup.py @@ -0,0 +1,90 @@ +from collections.abc import Callable +from time import monotonic, sleep +from typing import Final, Protocol + +from pydantic import BaseModel + +from batch_client import BatchObject, FileDeleteResponse +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError + +CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 600.0 +BATCH_CANCEL_POLL_SECONDS: Final = 10.0 + + +class BatchCleanupClient(Protocol): + def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ... + + def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ... + + def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ... + + +def cleanup_result[R: BaseModel]( + action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep +) -> Result[R]: + for delay, result in ((delay, action()) for delay in CLEANUP_DELAYS): + match result: + case NetworkError() | RateLimitedError(): + wait(delay) + case UnknownApiError(status_code=code) if code in {408, 429, 500, 502, 503, 504}: + wait(delay) + case _: + return result + return action() + + +def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) -> R: + match result: + case Success(data=data): + return data + case UnknownApiError(status_code=code): + raise AssertionError(f"{operation} failed: HTTP {code}") + case _: + raise AssertionError(f"{operation} failed: {result.kind}") + + +def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None: + result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider)) + if isinstance(result, UnknownApiError) and result.status_code == 404: + return + deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") + assert deleted.deleted, f"Delete file {file_id} did not confirm deletion" + + +def cleanup_batch( + client: BatchCleanupClient, + batch_id: str, + *, + key: str, + provider: str | None = None, + wait: Callable[[float], None] = sleep, + clock: Callable[[], float] = monotonic, +) -> None: + fetched: Final = _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} for cleanup", + ) + if fetched.status in BATCH_TERMINAL_STATUSES: + return + if fetched.status != "cancelling": + result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) + if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): + cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") + assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( + f"Cancel batch {batch_id} left status {cancelled.status}" + ) + deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS + while True: + current = _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} after cancellation", + ) + if current.status in BATCH_TERMINAL_STATUSES: + return + assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" + assert clock() < deadline, ( + f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" + ) + wait(BATCH_CANCEL_POLL_SECONDS) diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 31e49f22450..84a902b0b11 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -13,8 +13,9 @@ co-located here because only this suite uses them. from __future__ import annotations from dataclasses import dataclass +from typing import Final, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field from proxy_client import ProxyClient from e2e_http import ( @@ -27,6 +28,18 @@ from e2e_http import ( from models import LiteLLMParamsBody UPLOAD_FILENAME = "batch_input.jsonl" +AZURE_FILE_EXPIRY_SECONDS: Final = 14 * 24 * 60 * 60 + + +class ExpiringFileUploadForm(FileUploadForm): + expires_after_anchor: Literal["created_at"] = Field(default="created_at", alias="expires_after[anchor]") + expires_after_seconds: int = Field(default=AZURE_FILE_EXPIRY_SECONDS, alias="expires_after[seconds]") + + +def batch_upload_form(provider: str, *, target_model_names: str | None = None) -> FileUploadForm: + if provider == "azure": + return ExpiringFileUploadForm(target_model_names=target_model_names) + return FileUploadForm(target_model_names=target_model_names) class FileObject(BaseModel): @@ -37,6 +50,7 @@ class FileObject(BaseModel): bytes: int | None = None status: str | None = None created_at: int | None = None + expires_at: int | None = None class FileList(BaseModel): diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 1bcea0a61ee..17749c2fb87 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -108,6 +108,10 @@ class Capability: def id(self) -> str: return f"{self.provider}-{self.scenario}" + @property + def file_provider(self) -> str | None: + return self.provider if self.scenario in {"model_param", "provider_fallback"} else None + @property def jsonl_model(self) -> str: # Always the provider deployment name. Unified routes via diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 3b133fab680..91a365b6b92 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -13,7 +13,7 @@ the proxy config. from __future__ import annotations import os -from typing import Iterator +from typing import Final, Iterator import pytest @@ -21,6 +21,7 @@ from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_config import MANAGED_FILES_OPT_IN_ENV from e2e_http import NoBody +from lifecycle import ResourceManager from proxy_client import ProxyClient @@ -52,6 +53,13 @@ def client(proxy: ProxyClient) -> BatchClient: return build_client(proxy) +@pytest.fixture +def resources(client: BatchClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + yield manager + manager.teardown() + + @pytest.fixture(scope="session") def batch_deployments(client: BatchClient) -> Iterator[None]: probe = client.proxy.probe("/health/liveliness", params=NoBody()) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py new file mode 100644 index 00000000000..8ebfd750360 --- /dev/null +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -0,0 +1,187 @@ +from builtins import ExceptionGroup +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Final + +import pytest + +from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result +from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form +from capabilities import CAPABILITIES, Capability +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from lifecycle import ResourceManager +from models import KeyGenerateBody + + +@dataclass +class CleanupClient: + files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) + batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + calls: list[str] = field(default_factory=list) + + def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: + self.calls.append(f"delete {provider} {file_id}") + return next(self.files) + + def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls.append(f"retrieve {provider} {batch_id}") + return next(self.batches) + + def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls.append(f"cancel {provider} {batch_id}") + return next(self.cancellations) + + def generate_key(self, body: KeyGenerateBody) -> str: + return "test-key" + + def delete_key(self, key: str) -> None: + self.calls.append(f"delete key {key}") + + def delete_customers(self, user_ids: list[str]) -> None: + self.calls.append(f"delete customers {user_ids}") + + +def batch(status: str) -> Success[BatchObject]: + return Success(status_code=200, data=BatchObject(id="batch-1", status=status)) + + +def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]: + return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted)) + + +class TestFileCleanup: + @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) + def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: + client: Final = CleanupClient(files=iter((deleted_file(),))) + cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) + expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None + assert client.calls == [f"delete {expected_provider} file-1"] + + def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="secret response"),))) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure")) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + assert client.calls == ["delete azure file-1", "delete key test-key"] + assert len(caught.value.exceptions) == 1 + assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403" + + def test_success_response_must_confirm_deletion(self) -> None: + client: Final = CleanupClient(files=iter((deleted_file(deleted=False),))) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, "file-1", key="test-key") + + def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=404, body="missing"),))) + cleanup_file(client, "file-1", key="test-key", provider="azure") + assert client.calls == ["delete azure file-1"] + + def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="forbidden"),))) + manager: Final = ResourceManager(client=client) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.teardown() + assert client.calls == ["delete None file-1", "delete key test-key"] + + +class TestCleanupRetries: + @pytest.mark.parametrize( + "failure", + [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")], + ) + def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: + outcomes: Final = iter((failure, deleted_file())) + delays: Final[list[float]] = [] + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + assert isinstance(result, Success) and result.data.deleted + assert delays == [1.0] + + def test_persistent_error_has_bounded_retries(self) -> None: + failure: Final = UnknownApiError(status_code=503, body="unavailable") + outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) + delays: Final[list[float]] = [] + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + assert result is failure + assert tuple(delays) == CLEANUP_DELAYS + assert next(outcomes, None) is None + + def test_permanent_error_is_not_retried(self) -> None: + failure: Final = UnknownApiError(status_code=403, body="forbidden") + outcomes: Final = iter((failure, deleted_file())) + delays: Final[list[float]] = [] + assert cleanup_result(lambda: next(outcomes), wait=delays.append) is failure + assert delays == [] + assert isinstance(next(outcomes), Success) + + +class TestBatchCancellation: + def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: + client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) + delays: Final[list[float]] = [] + cleanup_batch(client, "batch-1", key="test-key", wait=delays.append) + assert client.calls == ["retrieve None batch-1"] * 3 + assert delays == [10.0] + + def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: + client: Final = CleanupClient( + batches=iter((batch("cancelling"), batch("cancelling"))), files=iter((deleted_file(),)) + ) + ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.defer(lambda: cleanup_batch(client, "batch-1", key=key, clock=lambda: next(ticks))) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + assert "cancellation did not finish" in str(caught.value.exceptions[0]) + assert client.calls == [ + "retrieve None batch-1", + "retrieve None batch-1", + "delete None file-1", + "delete key test-key", + ] + + @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) + def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: + client: Final = CleanupClient(batches=iter((batch(status),))) + cleanup_batch(client, "batch-1", key="test-key") + assert client.calls == ["retrieve None batch-1"] + + def test_active_batch_is_cancelled_through_its_provider(self) -> None: + client: Final = CleanupClient( + batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) + ) + cleanup_batch(client, "batch-1", key="test-key", provider="azure") + assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1", "retrieve azure batch-1"] + + @pytest.mark.parametrize("status", ["completed", "in_progress"]) + def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: + client: Final = CleanupClient( + batches=iter((batch("in_progress"), batch(status))), + cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), + ) + if status == "completed": + cleanup_batch(client, "batch-1", key="test-key") + else: + with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"): + cleanup_batch(client, "batch-1", key="test-key") + assert client.calls == ["retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"] + + +class TestAzureFileExpiry: + def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None: + form: Final = batch_upload_form("azure", target_model_names="azure-test") + assert form.model_dump(by_alias=True, exclude_none=True) == { + "purpose": "batch", + "target_model_names": "azure-test", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS, + } + + @pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"]) + def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None: + assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"} diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index ed7cf656d01..1b4a6ed266f 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -21,14 +21,16 @@ import os import re import time from datetime import datetime, timedelta, timezone -from typing import Callable import pytest from pydantic import BaseModel from e2e_config import PROXY_BASE_URL, unique_marker +from batch_cleanup import cleanup_batch, cleanup_file from batch_client import ( + AZURE_FILE_EXPIRY_SECONDS, + batch_upload_form, UPLOAD_FILENAME, BatchClient, BatchCreateBody, @@ -155,19 +157,19 @@ def upload_for_scenario( if cap.scenario == "encoded": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), model=cap.model, key=key, ) if cap.scenario == "unified": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch", target_model_names=cap.model), + form=batch_upload_form(cap.provider, target_model_names=cap.model), key=key, ) return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), key=key, provider=cap.provider, ) @@ -188,20 +190,11 @@ def create_for_scenario( def op_provider(cap: Capability) -> str | None: - """provider_fallback ids are raw, so retrieve/cancel/list/delete need the provider + """provider_fallback batch ids are raw, so retrieve/cancel/list need the provider hint; the other scenarios encode it into the id and route automatically.""" return cap.provider if cap.scenario == "provider_fallback" else None -def quietly(action: Callable[[], object]) -> Callable[[], None]: - """Adapt a value-returning call into a best-effort cleanup the teardown can run.""" - - def run() -> None: - action() - - return run - - def assert_file_object(file: FileObject, *, provider: str) -> None: assert file.object == "file", f"file.object={file.object!r}" assert file.purpose == "batch", f"file.purpose={file.purpose!r}" @@ -209,6 +202,10 @@ def assert_file_object(file: FileObject, *, provider: str) -> None: if provider != "bedrock": assert file.bytes > 0, f"file.bytes={file.bytes!r}" assert file.status, "file.status missing" + if provider == "azure": + assert file.expires_at is not None, "Azure batch input has no automatic expiry" + assert file.created_at is not None + assert file.expires_at - file.created_at == AZURE_FILE_EXPIRY_SECONDS assert ( file.created_at is not None and file.created_at > 0 ), "file.created_at missing" @@ -249,7 +246,7 @@ def test_batch_lifecycle( file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key)) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) + lambda: cleanup_file(client, file.id, key=key, provider=cap.file_provider) ) assert_file_object(file, provider=cap.provider) assert matches_id_shape( @@ -260,7 +257,7 @@ def test_batch_lifecycle( require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( - quietly(lambda: client.cancel_batch(batch.id, key=key, provider=provider)) + lambda: cleanup_batch(client, batch.id, key=key, provider=provider) ) assert batch.id, f"create returned no batch id (body={created.body[:200]})" @@ -339,7 +336,7 @@ def test_batch_key_model_access_denied( denied_upload = client.upload_file( content=render_jsonl(AZURE_BATCH_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) @@ -356,7 +353,7 @@ def test_batch_key_model_access_denied( ) ).id resources.defer( - quietly(lambda: client.delete_file(raw_file, key=key, provider="openai")) + lambda: cleanup_file(client, raw_file, key=key, provider="openai") ) denied_create = client.create_batch( @@ -383,6 +380,7 @@ def test_file_upload_and_delete_outputs( key=key, ) ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="openai") deleted = unwrap(client.delete_file(file.id, key=key)) @@ -458,12 +456,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) _ = client.proxy.poll_logs_for_key(key, min_rows=1) @@ -517,7 +515,7 @@ class TestBatchFileContent: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert file.id downloaded = client.proxy.transport.download( @@ -559,11 +557,11 @@ class TestBatchFileContent: file = unwrap( client.upload_file( content=payload, - form=FileUploadForm(purpose="batch", target_model_names=provider.model), + form=batch_upload_form(provider.name, target_model_names=provider.model), key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider=provider.name) assert is_managed_id(file.id), ( f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" @@ -626,7 +624,7 @@ class TestOpenAIFiles: ) ) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider="openai")) + lambda: cleanup_file(client, file.id, key=key, provider="openai") ) listed = unwrap(client.list_files(key=key)) @@ -690,7 +688,7 @@ class TestOpenAIFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) fetched = unwrap(client.retrieve_file(file.id, key=key)) assert fetched.id == file.id, "retrieve must echo the uploaded file id" @@ -760,7 +758,7 @@ class TestBatchRateLimitErrorMapping: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -813,7 +811,7 @@ class TestBatchEnqueuedTokenLimit: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) return file def _generate_enqueued_key( @@ -861,7 +859,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) @pytest.mark.covers( "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", @@ -904,7 +902,7 @@ class TestBatchEnqueuedTokenLimit: first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(first) first_batch = BatchObject.model_validate_json(first.body) - resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, first_batch.id, key=key)) blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) assert blocked.status_code == 429, ( @@ -928,7 +926,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(retried) retry_batch = BatchObject.model_validate_json(retried.body) - resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, retry_batch.id, key=key)) ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -984,13 +982,13 @@ class TestBedrockBatchAssumeRole: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="bedrock") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"assume-role create returned no batch id: {created.body[:200]}" assert is_managed_id(batch.id), ( @@ -1044,7 +1042,7 @@ class TestGeminiFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="gemini") assert file.id, "gemini file upload returned no id" @@ -1099,13 +1097,13 @@ class TestHostedVllmBatch: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="hosted_vllm") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}" assert batch.status in CREATED_BATCH_STATUSES, ( @@ -1192,7 +1190,7 @@ class TestBatchFailurePaths: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) @@ -1243,12 +1241,12 @@ class TestBatchFailurePaths: file = unwrap( client.upload_file( content=render_jsonl(AZURE_BATCH_RAW_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( f"upload did not encode the azure deployment into the file id: {file.id!r}" ) @@ -1258,7 +1256,7 @@ class TestBatchFailurePaths: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( "create with a foreign encoded file id must route by the file's embedded model, " @@ -1307,7 +1305,7 @@ class TestBatchSecondHop: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert is_managed_id(file.id), ( f"second-hop unified upload must return a managed file id, got {file.id!r}" ) @@ -1315,7 +1313,7 @@ class TestBatchSecondHop: created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert is_managed_id(batch.id), ( f"second-hop create must return a managed batch id, got {batch.id!r}" diff --git a/tests/e2e/batches/test_managed_files_enforcement_e2e.py b/tests/e2e/batches/test_managed_files_enforcement_e2e.py index 7ad0b16adc3..4f703cf0fdc 100644 --- a/tests/e2e/batches/test_managed_files_enforcement_e2e.py +++ b/tests/e2e/batches/test_managed_files_enforcement_e2e.py @@ -21,6 +21,7 @@ from typing import Iterator import pytest from batch_client import BatchClient, FileObject +from batch_cleanup import cleanup_file from capabilities import batch_model_name, is_managed_id, openai_batch_params from e2e_config import unique_marker from e2e_http import FileUploadForm, Result, UnknownApiError, unwrap @@ -108,7 +109,7 @@ def test_cross_user_managed_id_denied_owner_allowed( key=owner_key, ) ) - resources.defer(lambda: client.delete_file(uploaded.id, key=owner_key)) + resources.defer(lambda: cleanup_file(client, uploaded.id, key=owner_key)) assert is_managed_id(uploaded.id), f"expected a managed unified file id, got {uploaded.id}" denied = client.retrieve_file(uploaded.id, key=other_key) diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index c9a67ebdb8c..eb9704d4dcb 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -8,8 +8,9 @@ ResourceManager; the test registers a cleanup for every resource it creates, and the fixture's teardown releases them all even when the test body raises. """ +from builtins import ExceptionGroup from dataclasses import dataclass, field -from typing import Callable, List, Protocol, runtime_checkable +from typing import Callable, Final, List, Protocol, runtime_checkable from proxy_client import ProxyClient from models import KeyGenerateBody @@ -52,6 +53,7 @@ class ResourceManager: """ client: ResourceClient + strict_cleanup: bool = False _cleanups: List[Callable[[], object]] = field( default_factory=list ) # mutable-ok: append-only teardown registry @@ -82,8 +84,17 @@ class ResourceManager: return customer_id def teardown(self) -> None: - for cleanup in reversed(self._cleanups): - try: - cleanup() - except Exception: - pass # best-effort: a failed cleanup must not block the rest + failures: Final = tuple( + failure for cleanup in reversed(self._cleanups) + if (failure := _run_cleanup(cleanup)) is not None + ) + if failures and self.strict_cleanup: + raise ExceptionGroup("Resource cleanup failed", failures) + + +def _run_cleanup(cleanup: Callable[[], object]) -> Exception | None: + try: + cleanup() + except Exception as exc: + return exc + return None From a56c60e8924f6f956200c816142fb0cb994fa3ed Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:27:51 -0700 Subject: [PATCH 011/157] fix(e2e): wait for managed batch cancellation before deleting inputs --- tests/e2e/batches/COVERAGE.md | 5 +++-- tests/e2e/batches/batch_cleanup.py | 10 +++++++++- tests/e2e/batches/test_batch_cleanup.py | 14 ++++++++------ 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 899530dde2b..69ba9d781ec 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -129,8 +129,9 @@ provider when deleted. Model-encoded and managed file IDs route themselves File deletion and batch cancellation check their responses and retry transient failures up to three times. Teardown attempts every registered cleanup before reporting failures as test errors. Already deleted files and batches that are -terminal are safe to clean up again. Cancellation polls for up to ten minutes -before input deletion, because accepting cancellation does not finish it +terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes +before input deletion: the ten-minute provider window plus a propagation margin. +Raw and model-encoded inputs can be deleted after cancellation is accepted Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index a5b5e9bba37..dd79c776758 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -5,11 +5,12 @@ from typing import Final, Protocol from pydantic import BaseModel from batch_client import BatchObject, FileDeleteResponse +from capabilities import is_managed_id from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) -BATCH_CANCEL_TIMEOUT_SECONDS: Final = 600.0 +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 BATCH_CANCEL_POLL_SECONDS: Final = 10.0 @@ -62,12 +63,15 @@ def cleanup_batch( wait: Callable[[float], None] = sleep, clock: Callable[[], float] = monotonic, ) -> None: + needs_terminal_state: Final = is_managed_id(batch_id) fetched: Final = _require_cleanup_success( cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), f"Retrieve batch {batch_id} for cleanup", ) if fetched.status in BATCH_TERMINAL_STATUSES: return + if fetched.status == "cancelling" and not needs_terminal_state: + return if fetched.status != "cancelling": result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): @@ -75,6 +79,8 @@ def cleanup_batch( assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( f"Cancel batch {batch_id} left status {cancelled.status}" ) + if not needs_terminal_state: + return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS while True: current = _require_cleanup_success( @@ -84,6 +90,8 @@ def cleanup_batch( if current.status in BATCH_TERMINAL_STATUSES: return assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" + if not needs_terminal_state: + return assert clock() < deadline, ( f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" ) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 8ebfd750360..9715370d9b7 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -12,6 +12,8 @@ from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApi from lifecycle import ResourceManager from models import KeyGenerateBody +MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" + @dataclass class CleanupClient: @@ -122,8 +124,8 @@ class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) delays: Final[list[float]] = [] - cleanup_batch(client, "batch-1", key="test-key", wait=delays.append) - assert client.calls == ["retrieve None batch-1"] * 3 + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays.append) + assert client.calls == [f"retrieve None {MANAGED_BATCH_ID}"] * 3 assert delays == [10.0] def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: @@ -134,13 +136,13 @@ class TestBatchCancellation: manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) - manager.defer(lambda: cleanup_batch(client, "batch-1", key=key, clock=lambda: next(ticks))) + manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=lambda: next(ticks))) with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) assert client.calls == [ - "retrieve None batch-1", - "retrieve None batch-1", + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", "delete None file-1", "delete key test-key", ] @@ -156,7 +158,7 @@ class TestBatchCancellation: batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") - assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1", "retrieve azure batch-1"] + assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1"] @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: From 6bdf206cc5228e85b47c908a2f4977374656289b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:48:07 -0700 Subject: [PATCH 012/157] fix(e2e): accept managed file deletion responses --- tests/e2e/batches/COVERAGE.md | 3 +++ tests/e2e/batches/batch_cleanup.py | 4 +++- tests/e2e/batches/batch_client.py | 2 +- tests/e2e/batches/test_batch_cleanup.py | 15 +++++++++++++++ 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 69ba9d781ec..f95ea1f2649 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -139,6 +139,9 @@ fallback for interrupted runs: immediate deletion remains the normal cleanup. Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot be requested through its Files API +The Azure entry in `files_settings` must use `api_version: 2025-04-01-preview` +for raw uploads to honor expiry, matching the batch deployment's API version + ## Terminal state + cost write-back (cross-run marker baton) The 24h completion window rules out submit-and-wait inside one run, so diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index dd79c776758..3fd6802f696 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -51,7 +51,9 @@ def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider if isinstance(result, UnknownApiError) and result.status_code == 404: return deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") - assert deleted.deleted, f"Delete file {file_id} did not confirm deletion" + assert deleted.deleted is True or ( + deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file" + ), f"Delete file {file_id} did not confirm deletion" def cleanup_batch( diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 84a902b0b11..c9c77e1f12e 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -99,7 +99,7 @@ class BatchList(BaseModel): class FileDeleteResponse(BaseModel): id: str object: str | None = None - deleted: bool + deleted: bool | None = None class BatchCreateBody(BaseModel): diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 9715370d9b7..15dead6d36d 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -12,6 +12,7 @@ from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApi from lifecycle import ResourceManager from models import KeyGenerateBody +MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" @@ -53,6 +54,20 @@ def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]: class TestFileCleanup: + def test_managed_delete_accepts_the_deleted_file_object(self) -> None: + response: Final = Success( + status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) + ) + client: Final = CleanupClient(files=iter((response,))) + cleanup_file(client, MANAGED_FILE_ID, key="test-key") + assert client.calls == [f"delete None {MANAGED_FILE_ID}"] + + @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) + def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: + client: Final = CleanupClient(files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),))) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, file_id, key="test-key") + @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: client: Final = CleanupClient(files=iter((deleted_file(),))) From a096dd615c71e40be9473338ffeb8d1c17284f51 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 14:39:25 -0700 Subject: [PATCH 013/157] refactor(e2e): use immutable batch cleanup test expectations --- tests/e2e/batches/batch_cleanup.py | 7 +- tests/e2e/batches/test_batch_cleanup.py | 137 ++++++++++++++++-------- 2 files changed, 96 insertions(+), 48 deletions(-) diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 3fd6802f696..722df0c29bc 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -1,4 +1,5 @@ from collections.abc import Callable +from itertools import count from time import monotonic, sleep from typing import Final, Protocol @@ -84,11 +85,13 @@ def cleanup_batch( if not needs_terminal_state: return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS - while True: - current = _require_cleanup_success( + for current in ( + _require_cleanup_success( cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), f"Retrieve batch {batch_id} after cancellation", ) + for _ in count() + ): if current.status in BATCH_TERMINAL_STATUSES: return assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 15dead6d36d..66b7079ccc2 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -16,33 +16,44 @@ MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" -@dataclass +@dataclass(frozen=True, slots=True) +class ExpectedCalls[T]: + values: Iterator[T] + + def __call__(self, value: T) -> None: + assert next(self.values, None) == value + + def assert_done(self) -> None: + assert tuple(self.values) == () + + +@dataclass(frozen=True, slots=True) class CleanupClient: + calls: ExpectedCalls[str] files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) - calls: list[str] = field(default_factory=list) def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: - self.calls.append(f"delete {provider} {file_id}") + self.calls(f"delete {provider} {file_id}") return next(self.files) def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: - self.calls.append(f"retrieve {provider} {batch_id}") + self.calls(f"retrieve {provider} {batch_id}") return next(self.batches) def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: - self.calls.append(f"cancel {provider} {batch_id}") + self.calls(f"cancel {provider} {batch_id}") return next(self.cancellations) def generate_key(self, body: KeyGenerateBody) -> str: return "test-key" def delete_key(self, key: str) -> None: - self.calls.append(f"delete key {key}") + self.calls(f"delete key {key}") def delete_customers(self, user_ids: list[str]) -> None: - self.calls.append(f"delete customers {user_ids}") + self.calls(f"delete customers {user_ids}") def batch(status: str) -> Success[BatchObject]: @@ -58,51 +69,71 @@ class TestFileCleanup: response: Final = Success( status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) ) - client: Final = CleanupClient(files=iter((response,))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete None {MANAGED_FILE_ID}",))), files=iter((response,)) + ) cleanup_file(client, MANAGED_FILE_ID, key="test-key") - assert client.calls == [f"delete None {MANAGED_FILE_ID}"] + client.calls.assert_done() @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: - client: Final = CleanupClient(files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete None {file_id}",))), + files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),)), + ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, file_id, key="test-key") + client.calls.assert_done() @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: - client: Final = CleanupClient(files=iter((deleted_file(),))) - cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None - assert client.calls == [f"delete {expected_provider} file-1"] + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete {expected_provider} file-1",))), files=iter((deleted_file(),)) + ) + cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) + client.calls.assert_done() def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="secret response"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete azure file-1", "delete key test-key"))), + files=iter((UnknownApiError(status_code=403, body="secret response"),)), + ) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure")) with pytest.raises(ExceptionGroup) as caught: manager.teardown() - assert client.calls == ["delete azure file-1", "delete key test-key"] + client.calls.assert_done() assert len(caught.value.exceptions) == 1 assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403" def test_success_response_must_confirm_deletion(self) -> None: - client: Final = CleanupClient(files=iter((deleted_file(deleted=False),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete None file-1",))), files=iter((deleted_file(deleted=False),)) + ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, "file-1", key="test-key") + client.calls.assert_done() def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=404, body="missing"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete azure file-1",))), + files=iter((UnknownApiError(status_code=404, body="missing"),)), + ) cleanup_file(client, "file-1", key="test-key", provider="azure") - assert client.calls == ["delete azure file-1"] + client.calls.assert_done() def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="forbidden"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete None file-1", "delete key test-key"))), + files=iter((UnknownApiError(status_code=403, body="forbidden"),)), + ) manager: Final = ResourceManager(client=client) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) manager.teardown() - assert client.calls == ["delete None file-1", "delete key test-key"] + client.calls.assert_done() class TestCleanupRetries: @@ -112,40 +143,54 @@ class TestCleanupRetries: ) def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: outcomes: Final = iter((failure, deleted_file())) - delays: Final[list[float]] = [] - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + delays: Final = ExpectedCalls(iter((1.0,))) + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) assert isinstance(result, Success) and result.data.deleted - assert delays == [1.0] + delays.assert_done() def test_persistent_error_has_bounded_retries(self) -> None: failure: Final = UnknownApiError(status_code=503, body="unavailable") outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) - delays: Final[list[float]] = [] - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + delays: Final = ExpectedCalls(iter(CLEANUP_DELAYS)) + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) assert result is failure - assert tuple(delays) == CLEANUP_DELAYS + delays.assert_done() assert next(outcomes, None) is None def test_permanent_error_is_not_retried(self) -> None: failure: Final = UnknownApiError(status_code=403, body="forbidden") outcomes: Final = iter((failure, deleted_file())) - delays: Final[list[float]] = [] - assert cleanup_result(lambda: next(outcomes), wait=delays.append) is failure - assert delays == [] + delays: Final = ExpectedCalls[float](iter(())) + assert cleanup_result(lambda: next(outcomes), wait=delays) is failure + delays.assert_done() assert isinstance(next(outcomes), Success) class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: - client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) - delays: Final[list[float]] = [] - cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays.append) - assert client.calls == [f"retrieve None {MANAGED_BATCH_ID}"] * 3 - assert delays == [10.0] + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"retrieve None {MANAGED_BATCH_ID}",) * 3)), + batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled"))), + ) + delays: Final = ExpectedCalls(iter((10.0,))) + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays) + client.calls.assert_done() + delays.assert_done() def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: client: Final = CleanupClient( - batches=iter((batch("cancelling"), batch("cancelling"))), files=iter((deleted_file(),)) + calls=ExpectedCalls( + iter( + ( + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", + "delete None file-1", + "delete key test-key", + ) + ) + ), + batches=iter((batch("cancelling"), batch("cancelling"))), + files=iter((deleted_file(),)), ) ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) manager: Final = ResourceManager(client=client, strict_cleanup=True) @@ -155,29 +200,29 @@ class TestBatchCancellation: with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) - assert client.calls == [ - f"retrieve None {MANAGED_BATCH_ID}", - f"retrieve None {MANAGED_BATCH_ID}", - "delete None file-1", - "delete key test-key", - ] + client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: - client: Final = CleanupClient(batches=iter((batch(status),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("retrieve None batch-1",))), batches=iter((batch(status),)) + ) cleanup_batch(client, "batch-1", key="test-key") - assert client.calls == ["retrieve None batch-1"] + client.calls.assert_done() def test_active_batch_is_cancelled_through_its_provider(self) -> None: client: Final = CleanupClient( - batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) + calls=ExpectedCalls(iter(("retrieve azure batch-1", "cancel azure batch-1"))), + batches=iter((batch("in_progress"), batch("cancelled"))), + cancellations=iter((batch("cancelling"),)), ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") - assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1"] + client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( + calls=ExpectedCalls(iter(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"))), batches=iter((batch("in_progress"), batch(status))), cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), ) @@ -186,7 +231,7 @@ class TestBatchCancellation: else: with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"): cleanup_batch(client, "batch-1", key="test-key") - assert client.calls == ["retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"] + client.calls.assert_done() class TestAzureFileExpiry: From c84131b81ab6feb552e23c5996559159e8810066 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 7 Sep 2026 16:00:20 -0700 Subject: [PATCH 014/157] feat(proxy): price cache and reasoning tokens in /cost/estimate POST /cost/estimate now accepts cache_read_input_tokens, cache_creation_input_tokens and reasoning_tokens, bills them at the model's cache and reasoning rates, and reports each share per request, per day and per month next to the rates it used. Custom-priced deployments also get cache and reasoning lines in the cost breakdown now, so the estimate and the spend logs reconcile with their totals instead of showing zero for those tokens. Requested by a customer (Pylon #7365). Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- litellm/cost_calculator.py | 1 + .../litellm_core_utils/llm_cost_calc/utils.py | 68 +++-- litellm/proxy/_types.py | 36 +++ .../cost_tracking_settings.py | 253 +++++++++++------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 64 +++++ .../test_cost_tracking_settings.py | 187 ++++++++++++- tests/test_litellm/test_cost_calculator.py | 50 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 103 ++++++- 8 files changed, 641 insertions(+), 121 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9a9d2ceda03..a1181ee4124 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1746,6 +1746,7 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + custom_cost_per_token=custom_cost_per_token, ) _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 68dc27ec25e..8d53c8da3e6 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -19,6 +19,7 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, CompletionTokensDetailsWrapper, + CostPerToken, DataResidency, ImageResponse, ModelInfo, @@ -1307,6 +1308,40 @@ class TokenTypeCostBreakdown: cache_creation_cost: float +def _reasoning_token_count(usage: Usage) -> int: + parsed: Final = ( + parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 + ) + return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) + + +def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details + first, then the private top-level counters the Usage constructor mirrors cache tokens onto for + providers/callers that bypass the details.""" + parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None + parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 + parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 + return ( + parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), + parsed["cache_creation_token_details"] if parsed is not None else None, + ) + + +def _custom_pricing_token_type_breakdown(usage: Usage, custom_cost_per_token: CostPerToken) -> TokenTypeCostBreakdown: + """Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured + cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" + input_rate: Final = custom_cost_per_token["input_cost_per_token"] + cache_read_tokens, cache_creation_tokens, _ = _cache_token_counts(usage) + return TokenTypeCostBreakdown( + reasoning_cost=float(_reasoning_token_count(usage)) * custom_cost_per_token["output_cost_per_token"], + cache_read_cost=float(cache_read_tokens) * custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_creation_cost=float(cache_creation_tokens) + * custom_cost_per_token.get("cache_creation_input_token_cost", input_rate), + ) + + def get_token_type_cost_breakdown( model: str, custom_llm_provider: str | None, @@ -1315,6 +1350,7 @@ def get_token_type_cost_breakdown( data_residency: str | None = None, vertex_location: str | None = None, current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, ) -> TokenTypeCostBreakdown: """ Provider-agnostic cost of reasoning and cache tokens, derived from the usage @@ -1325,9 +1361,13 @@ def get_token_type_cost_breakdown( land on ``prompt_tokens_details`` (via the Usage constructor and provider transformations) and reasoning tokens on ``completion_tokens_details``. It reuses the same rate-resolution primitives as the total-cost path so the breakdown can - never drift from the totals. Returns zeros (never raises) when the model or its - pricing cannot be resolved. + never drift from the totals. A deployment billed by ``custom_cost_per_token`` is + priced from those flat rates instead of the cost map, for the same reason. + Returns zeros (never raises) when the model or its pricing cannot be resolved. """ + if custom_cost_per_token is not None: + return _custom_pricing_token_type_breakdown(usage=usage, custom_cost_per_token=custom_cost_per_token) + try: model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: @@ -1348,12 +1388,6 @@ def get_token_type_cost_breakdown( threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) - reasoning_tokens = ( - parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 - ) - if not reasoning_tokens: - reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - reasoning_rate: Final = _resolve_billed_reasoning_rate( model_info=model_info, usage=usage, @@ -1361,23 +1395,9 @@ def get_token_type_cost_breakdown( completion_base_cost=completion_base_cost, current_time=billing_time, ) - reasoning_cost = float(reasoning_tokens) * reasoning_rate - - cache_read_tokens = 0 - cache_creation_tokens = 0 - cache_creation_token_details: CacheCreationTokenDetails | None = None - if usage.prompt_tokens_details is not None: - prompt_tokens_details: Final = parse_prompt_tokens_details(usage) - cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] - cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] - cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] - # Fall back to the private top-level counters the Usage constructor mirrors cache - # tokens onto, so providers/callers that bypass prompt_tokens_details are covered. - if not cache_read_tokens: - cache_read_tokens = _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)) - if not cache_creation_tokens: - cache_creation_tokens = _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)) + reasoning_cost = float(_reasoning_token_count(usage)) * reasoning_rate + cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate cache_creation_cost = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index abce10690e5..65b7d409d7c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5137,9 +5137,26 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase): model: str = Field(description="Model name (from /model_group/info)") input_tokens: int = Field(description="Expected input tokens per request", ge=0) output_tokens: int = Field(description="Expected output tokens per request", ge=0) + cache_read_input_tokens: int = Field( + default=0, description="Input tokens read from the prompt cache; counted within input_tokens", ge=0 + ) + cache_creation_input_tokens: int = Field( + default=0, description="Input tokens written to the prompt cache; counted within input_tokens", ge=0 + ) + reasoning_tokens: int = Field( + default=0, description="Reasoning tokens the model emits; counted within output_tokens", ge=0 + ) num_requests_per_day: int | None = Field(default=None, description="Number of requests per day", ge=0) num_requests_per_month: int | None = Field(default=None, description="Number of requests per month", ge=0) + @model_validator(mode="after") + def validate_token_subsets(self) -> "CostEstimateRequest": + if self.cache_read_input_tokens + self.cache_creation_input_tokens > self.input_tokens: + raise ValueError("cache_read_input_tokens plus cache_creation_input_tokens cannot exceed input_tokens") + if self.reasoning_tokens > self.output_tokens: + raise ValueError("reasoning_tokens cannot exceed output_tokens") + return self + class CostEstimateResponse(LiteLLMPydanticObjectBase): """Response body for /cost/estimate endpoint.""" @@ -5147,6 +5164,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): model: str input_tokens: int output_tokens: int + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + reasoning_tokens: int = 0 num_requests_per_day: int | None = None num_requests_per_month: int | None = None # Per-request costs @@ -5154,17 +5174,33 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): input_cost_per_request: float = Field(description="Input token cost per request (before margin)") output_cost_per_request: float = Field(description="Output token cost per request (before margin)") margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request") + cache_read_cost_per_request: float = Field(default=0.0, description="Cache-read share of input_cost_per_request") + cache_creation_cost_per_request: float = Field( + default=0.0, description="Cache-write share of input_cost_per_request" + ) + reasoning_cost_per_request: float = Field(default=0.0, description="Reasoning share of output_cost_per_request") # Daily costs (if num_requests_per_day provided) daily_cost: float | None = Field(default=None, description="Total daily cost (includes margin)") daily_input_cost: float | None = Field(default=None, description="Daily input token cost") daily_output_cost: float | None = Field(default=None, description="Daily output token cost") daily_margin_cost: float | None = Field(default=None, description="Daily margin/fee") + daily_cache_read_cost: float | None = Field(default=None, description="Cache-read share of daily_input_cost") + daily_cache_creation_cost: float | None = Field(default=None, description="Cache-write share of daily_input_cost") + daily_reasoning_cost: float | None = Field(default=None, description="Reasoning share of daily_output_cost") # Monthly costs (if num_requests_per_month provided) monthly_cost: float | None = Field(default=None, description="Total monthly cost (includes margin)") monthly_input_cost: float | None = Field(default=None, description="Monthly input token cost") monthly_output_cost: float | None = Field(default=None, description="Monthly output token cost") monthly_margin_cost: float | None = Field(default=None, description="Monthly margin/fee") + monthly_cache_read_cost: float | None = Field(default=None, description="Cache-read share of monthly_input_cost") + monthly_cache_creation_cost: float | None = Field( + default=None, description="Cache-write share of monthly_input_cost" + ) + monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost") # Pricing info input_cost_per_token: float | None = None output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token") + cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token") + output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token") provider: str | None = None diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 204051c3715..51f31a757cd 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -27,7 +27,15 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo +from litellm.types.utils import ( + CostBreakdown, + CostPerToken, + LlmProvidersSet, + ModelInfo, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) router: Final = APIRouter() @@ -46,13 +54,15 @@ def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> fl def _extract_custom_pricing( - litellm_params: Mapping[str, object], model_info: Mapping[str, object] + litellm_params: Mapping[str, object], model_info: Mapping[str, object], builtin: ModelInfo | None ) -> CostPerToken | None: """ Pull per-token pricing configured on a deployment so on-prem / self-hosted models (absent from the public cost map) still estimate a real cost. Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params`` - wins, matching the router's cost-map registration precedence. + wins, matching the router's cost-map registration precedence. Cache rates the + deployment leaves unset come from the backend model's built-in entry, then its + own input rate, again matching what the router registers for live billing. """ sources: Final = (litellm_params, model_info) input_price: Final = _configured_price("input_cost_per_token", sources) @@ -61,9 +71,15 @@ def _extract_custom_pricing( if input_price is None and output_price is None: return None + input_rate: Final = input_price or 0.0 + cache_sources: Final = sources if builtin is None else (*sources, builtin) + cache_read_price: Final = _configured_price("cache_read_input_token_cost", cache_sources) + cache_creation_price: Final = _configured_price("cache_creation_input_token_cost", cache_sources) return CostPerToken( - input_cost_per_token=input_price or 0.0, + input_cost_per_token=input_rate, output_cost_per_token=output_price or 0.0, + cache_read_input_token_cost=input_rate if cache_read_price is None else cache_read_price, + cache_creation_input_token_cost=input_rate if cache_creation_price is None else cache_creation_price, ) @@ -98,17 +114,14 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: model_info: Final = first_deployment.get("model_info", {}) custom_llm_provider: Final = litellm_params.get("custom_llm_provider") provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None - custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info) - - # Check base_model first (needed for Azure custom deployment names) + # base_model wins (needed for Azure custom deployment names) base_model: Final = model_info.get("base_model") or litellm_params.get("base_model") - if base_model: - verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model) - return ResolvedCostModel(str(base_model), provider, custom_cost_per_token) - - resolved_model: Final = litellm_params.get("model") + resolved_model: Final = base_model or litellm_params.get("model") if resolved_model: verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model) + custom_cost_per_token: Final = _extract_custom_pricing( + litellm_params, model_info, _lookup_model_info(str(resolved_model)) + ) return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token) except Exception as e: verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e) @@ -117,19 +130,97 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: return ResolvedCostModel(model, None, None) -def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost): - """ - Calculate costs for a given number of requests. +@dataclass(frozen=True, slots=True) +class CostLines: + """Cost of one request split the way the spend logs split it: the cache lines are + shares of input_cost and the reasoning line is a share of output_cost.""" - Returns tuple of (total_cost, input_cost, output_cost, margin_cost) or all None if num_requests is None/0. - """ - if not num_requests: - return None, None, None, None - return ( - cost_per_request * num_requests, - input_cost * num_requests, - output_cost * num_requests, - margin_cost * num_requests, + total_cost: float + input_cost: float + output_cost: float + margin_cost: float + cache_read_cost: float + cache_creation_cost: float + reasoning_cost: float + + def times(self, num_requests: int | None) -> "CostLines | None": + if not num_requests: + return None + return CostLines( + total_cost=self.total_cost * num_requests, + input_cost=self.input_cost * num_requests, + output_cost=self.output_cost * num_requests, + margin_cost=self.margin_cost * num_requests, + cache_read_cost=self.cache_read_cost * num_requests, + cache_creation_cost=self.cache_creation_cost * num_requests, + reasoning_cost=self.reasoning_cost * num_requests, + ) + + +def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) -> CostLines: + breakdown: Final = cost_breakdown if cost_breakdown is not None else CostBreakdown() + return CostLines( + total_cost=cost_per_request, + input_cost=breakdown.get("input_cost", 0.0), + output_cost=breakdown.get("output_cost", 0.0), + margin_cost=breakdown.get("margin_total_amount", 0.0), + cache_read_cost=breakdown.get("cache_read_cost", 0.0), + cache_creation_cost=breakdown.get("cache_creation_cost", 0.0), + reasoning_cost=breakdown.get("reasoning_cost", 0.0), + ) + + +@dataclass(frozen=True, slots=True) +class EffectiveTokenRates: + input_cost_per_token: float | None + output_cost_per_token: float | None + cache_read_input_token_cost: float | None + cache_creation_input_token_cost: float | None + output_cost_per_reasoning_token: float | None + + +def _custom_token_rates(custom_cost_per_token: CostPerToken) -> EffectiveTokenRates: + input_rate: Final = custom_cost_per_token["input_cost_per_token"] + output_rate: Final = custom_cost_per_token["output_cost_per_token"] + return EffectiveTokenRates( + input_cost_per_token=input_rate, + output_cost_per_token=output_rate, + cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_creation_input_token_cost=custom_cost_per_token.get("cache_creation_input_token_cost", input_rate), + output_cost_per_reasoning_token=output_rate, + ) + + +def _cost_map_token_rates(model_info: ModelInfo | None) -> EffectiveTokenRates: + """Base rates the cost calculator bills flat usage at: a cost-map model without a cache price + bills cache tokens at zero, and one without a reasoning price bills reasoning at the output rate.""" + if model_info is None: + return EffectiveTokenRates(None, None, None, None, None) + sources: Final = (model_info,) + output_rate: Final = _configured_price("output_cost_per_token", sources) + reasoning_rate: Final = _configured_price("output_cost_per_reasoning_token", sources) + return EffectiveTokenRates( + input_cost_per_token=_configured_price("input_cost_per_token", sources), + output_cost_per_token=output_rate, + cache_read_input_token_cost=_configured_price("cache_read_input_token_cost", sources) or 0.0, + cache_creation_input_token_cost=_configured_price("cache_creation_input_token_cost", sources) or 0.0, + output_cost_per_reasoning_token=output_rate if reasoning_rate is None else reasoning_rate, + ) + + +def _usage_for_estimate(request: CostEstimateRequest) -> Usage: + cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens + return Usage( + prompt_tokens=request.input_tokens, + completion_tokens=request.output_tokens, + total_tokens=request.input_tokens + request.output_tokens, + reasoning_tokens=request.reasoning_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=request.cache_read_input_tokens, + cache_creation_tokens=request.cache_creation_input_tokens, + ) + if cache_tokens + else None, ) @@ -530,11 +621,14 @@ async def estimate_cost( - model: Model name (e.g., "gpt-4", "claude-3-opus") - input_tokens: Expected input tokens per request - output_tokens: Expected output tokens per request + - cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional) + - cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional) + - reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional) - num_requests_per_day: Number of requests per day (optional) - num_requests_per_month: Number of requests per month (optional) Returns cost breakdown including: - - Per-request costs (input, output, margin) + - Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares) - Daily costs (if num_requests_per_day provided) - Monthly costs (if num_requests_per_month provided) @@ -543,14 +637,15 @@ async def estimate_cost( { "model": "gpt-4", "input_tokens": 1000, + "cache_read_input_tokens": 800, "output_tokens": 500, + "reasoning_tokens": 200, "num_requests_per_day": 100, "num_requests_per_month": 3000 } ``` """ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import ModelResponse, Usage # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') resolved: Final = _resolve_model_for_cost_lookup(request.model) @@ -559,15 +654,7 @@ async def estimate_cost( verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model) - # Create a mock response with usage for completion_cost - mock_response: Final = ModelResponse( - model=resolved_model, - usage=Usage( - prompt_tokens=request.input_tokens, - completion_tokens=request.output_tokens, - total_tokens=request.input_tokens + request.output_tokens, - ), - ) + mock_response: Final = ModelResponse(model=resolved_model, usage=_usage_for_estimate(request)) # Create a logging object to capture cost breakdown litellm_logging_obj: Final = LiteLLMLoggingObj( @@ -597,75 +684,53 @@ async def estimate_cost( }, ) - # Get cost breakdown from the logging object - cost_breakdown: Final = litellm_logging_obj.cost_breakdown - - input_cost: Final = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0 - output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 - margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 + per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown) + daily: Final = per_request.times(request.num_requests_per_day) + monthly: Final = per_request.times(request.num_requests_per_month) model_info: Final = _lookup_model_info(resolved_model) - mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None - mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None + rates: Final = ( + _custom_token_rates(resolved.custom_cost_per_token) + if resolved.custom_cost_per_token is not None + else _cost_map_token_rates(model_info) + ) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None - - input_cost_per_token: Final = ( - resolved.custom_cost_per_token["input_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_input_price - ) - output_cost_per_token: Final = ( - resolved.custom_cost_per_token["output_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_output_price - ) custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider - # Calculate daily and monthly costs - ( - daily_cost, - daily_input_cost, - daily_output_cost, - daily_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_day, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - ( - monthly_cost, - monthly_input_cost, - monthly_output_cost, - monthly_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - return CostEstimateResponse( model=request.model, input_tokens=request.input_tokens, output_tokens=request.output_tokens, + cache_read_input_tokens=request.cache_read_input_tokens, + cache_creation_input_tokens=request.cache_creation_input_tokens, + reasoning_tokens=request.reasoning_tokens, num_requests_per_day=request.num_requests_per_day, num_requests_per_month=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost_per_request=input_cost, - output_cost_per_request=output_cost, - margin_cost_per_request=margin_cost, - daily_cost=daily_cost, - daily_input_cost=daily_input_cost, - daily_output_cost=daily_output_cost, - daily_margin_cost=daily_margin_cost, - monthly_cost=monthly_cost, - monthly_input_cost=monthly_input_cost, - monthly_output_cost=monthly_output_cost, - monthly_margin_cost=monthly_margin_cost, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, + cost_per_request=per_request.total_cost, + input_cost_per_request=per_request.input_cost, + output_cost_per_request=per_request.output_cost, + margin_cost_per_request=per_request.margin_cost, + cache_read_cost_per_request=per_request.cache_read_cost, + cache_creation_cost_per_request=per_request.cache_creation_cost, + reasoning_cost_per_request=per_request.reasoning_cost, + daily_cost=daily.total_cost if daily is not None else None, + daily_input_cost=daily.input_cost if daily is not None else None, + daily_output_cost=daily.output_cost if daily is not None else None, + daily_margin_cost=daily.margin_cost if daily is not None else None, + daily_cache_read_cost=daily.cache_read_cost if daily is not None else None, + daily_cache_creation_cost=daily.cache_creation_cost if daily is not None else None, + daily_reasoning_cost=daily.reasoning_cost if daily is not None else None, + monthly_cost=monthly.total_cost if monthly is not None else None, + monthly_input_cost=monthly.input_cost if monthly is not None else None, + monthly_output_cost=monthly.output_cost if monthly is not None else None, + monthly_margin_cost=monthly.margin_cost if monthly is not None else None, + monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None, + monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None, + monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None, + input_cost_per_token=rates.input_cost_per_token, + output_cost_per_token=rates.output_cost_per_token, + cache_read_input_token_cost=rates.cache_read_input_token_cost, + cache_creation_input_token_cost=rates.cache_creation_input_token_cost, + output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token, provider=custom_llm_provider, ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 59f0938e338..7065003839b 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -3874,6 +3874,70 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_co assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost) +def _custom_priced_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + +def test_token_type_cost_breakdown_prices_custom_pricing_from_its_flat_rates(): + """ + A custom-priced deployment, usually absent from the cost map, used to get zero cache and + reasoning lines while its total already billed cache tokens at the custom cache rates. + The lines must come from the same flat rates: a configured cache rate, else the input + rate for cache tokens and the output rate for reasoning tokens. + """ + from litellm.types.utils import CostPerToken + + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=_custom_priced_usage(), + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7 + ), + ) + + assert breakdown.cache_read_cost == pytest.approx(800 * 1e-7) + assert breakdown.cache_creation_cost == pytest.approx(100 * 1e-6) + assert breakdown.reasoning_cost == pytest.approx(200 * 2e-6) + + +def test_token_type_cost_breakdown_reconciles_with_custom_pricing_totals(): + from litellm.cost_calculator import cost_per_token + from litellm.types.utils import CostPerToken + + usage = _custom_priced_usage() + custom_cost_per_token = CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ) + + prompt_cost, completion_cost = cost_per_token( + model="openai/onprem-model", + custom_llm_provider="openai", + prompt_tokens=1000, + completion_tokens=500, + usage_object=usage, + custom_cost_per_token=custom_cost_per_token, + ) + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=usage, + custom_cost_per_token=custom_cost_per_token, + ) + + assert 100 * 1e-6 + breakdown.cache_read_cost + breakdown.cache_creation_cost == pytest.approx(prompt_cost) + assert 300 * 2e-6 + breakdown.reasoning_cost == pytest.approx(completion_cost) + + def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index ec62cc47018..6d3ef2ffea9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -8,9 +8,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from pydantic import ValidationError import litellm +from litellm.proxy._types import CostEstimateRequest from litellm.proxy.management_endpoints.cost_tracking_settings import router from litellm.proxy.proxy_server import app @@ -789,13 +791,13 @@ INPUT_TOKENS = 1000 OUTPUT_TOKENS = 500 -def _router_pricing(**pricing: float) -> MagicMock: +def _router_pricing(model: str = AN_UNDERLYING_MODEL, **pricing: float) -> MagicMock: mock_router = MagicMock() mock_router.get_model_list.return_value = [ { "model_name": AN_ALIAS, "litellm_params": { - "model": AN_UNDERLYING_MODEL, + "model": model, "custom_llm_provider": "openai", **pricing, }, @@ -909,3 +911,184 @@ class TestEstimateCostPeriodTotals: assert response.cost_per_request == pytest.approx(0.0022) assert response.daily_margin_cost == pytest.approx(0.02) assert response.daily_cost == pytest.approx(0.22) + + +CACHE_READ_TOKENS = 800 +CACHE_CREATION_TOKENS = 100 +REASONING_TOKENS = 200 +TEXT_INPUT_TOKENS = INPUT_TOKENS - CACHE_READ_TOKENS - CACHE_CREATION_TOKENS +TEXT_OUTPUT_TOKENS = OUTPUT_TOKENS - REASONING_TOKENS + + +async def _estimate_with_cache_and_reasoning(mock_router: MagicMock | None, model: str = AN_ALIAS, **overrides: int): + return await _estimate( + mock_router, + model=model, + cache_read_input_tokens=CACHE_READ_TOKENS, + cache_creation_input_tokens=CACHE_CREATION_TOKENS, + reasoning_tokens=REASONING_TOKENS, + **overrides, + ) + + +class TestEstimateCostCacheAndReasoningTokens: + @pytest.mark.asyncio + async def test_a_mapped_model_bills_cache_and_reasoning_tokens_at_their_own_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 3e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 3.75e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 1e-5) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 3e-6 + CACHE_READ_TOKENS * 3e-7 + CACHE_CREATION_TOKENS * 3.75e-6 + ) + assert response.output_cost_per_request == pytest.approx(TEXT_OUTPUT_TOKENS * 15e-6 + REASONING_TOKENS * 1e-5) + assert response.cost_per_request == pytest.approx( + response.input_cost_per_request + response.output_cost_per_request + ) + assert response.daily_cache_read_cost == pytest.approx(10 * CACHE_READ_TOKENS * 3e-7) + assert response.daily_cache_creation_cost == pytest.approx(10 * CACHE_CREATION_TOKENS * 3.75e-6) + assert response.daily_reasoning_cost == pytest.approx(10 * REASONING_TOKENS * 1e-5) + assert response.monthly_cache_read_cost is None + assert response.cache_read_input_token_cost == pytest.approx(3e-7) + assert response.cache_creation_input_token_cost == pytest.approx(3.75e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(1e-5) + assert ( + response.cache_read_input_tokens, + response.cache_creation_input_tokens, + response.reasoning_tokens, + ) == (CACHE_READ_TOKENS, CACHE_CREATION_TOKENS, REASONING_TOKENS) + + @pytest.mark.asyncio + async def test_a_model_without_cache_or_reasoning_prices_estimates_what_the_proxy_bills(self, monkeypatch): + """The cost calculator bills cache tokens of a cost-map model without cache prices at zero + and its reasoning tokens at the output rate. The estimate reports those effective rates.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + {"input_cost_per_token": 5e-6, "output_cost_per_token": 6e-6, "litellm_provider": "openai", "mode": "chat"}, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL) + + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 6e-6) + assert response.input_cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6) + assert response.cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6 + OUTPUT_TOKENS * 6e-6) + assert response.cache_read_input_token_cost == 0.0 + assert response.cache_creation_input_token_cost == 0.0 + assert response.output_cost_per_reasoning_token == pytest.approx(6e-6) + + @pytest.mark.asyncio + async def test_a_request_without_cache_or_reasoning_tokens_estimates_as_before(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cost_per_request == pytest.approx(INPUT_TOKENS * 3e-6 + OUTPUT_TOKENS * 15e-6) + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == 0.0 + assert response.daily_cache_read_cost == 0.0 + assert response.daily_reasoning_cost == 0.0 + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_bills_cache_and_reasoning_tokens_from_its_flat_rates(self): + response = await _estimate_with_cache_and_reasoning( + _router_pricing(input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 1e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 1e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 2e-6) + assert response.cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 1e-7 + CACHE_CREATION_TOKENS * 1e-6 + OUTPUT_TOKENS * 2e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(1e-7) + assert response.cache_creation_input_token_cost == pytest.approx(1e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(2e-6) + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_of_a_mapped_model_inherits_its_built_in_cache_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 5e-6, + "output_cost_per_token": 6e-6, + "cache_read_input_token_cost": 5e-7, + "cache_creation_input_token_cost": 6.25e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning( + _router_pricing(model=A_MAPPED_MODEL, input_cost_per_token=1e-6, output_cost_per_token=2e-6) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 5e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 6.25e-6) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 5e-7 + CACHE_CREATION_TOKENS * 6.25e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(5e-7) + assert response.cache_creation_input_token_cost == pytest.approx(6.25e-6) + + +class TestCostEstimateRequestTokenSubsets: + def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed input_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + cache_read_input_tokens=INPUT_TOKENS, + cache_creation_input_tokens=1, + ) + + def test_reasoning_tokens_beyond_the_output_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed output_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + reasoning_tokens=OUTPUT_TOKENS + 1, + ) + + def test_the_endpoint_answers_422_when_cache_tokens_exceed_input_tokens(self): + response = client.post( + "/cost/estimate", + headers={"Authorization": "Bearer sk-1234"}, + json={"model": AN_ALIAS, "input_tokens": 1000, "output_tokens": 100, "cache_read_input_tokens": 8000}, + ) + + assert response.status_code == 422 + assert "cannot exceed input_tokens" in response.text diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f8fa2231597..910587d3dc5 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3736,6 +3736,56 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100 * 3e-08) +def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing(): + """ + A custom-priced deployment bills cache tokens at its custom cache rates, but the + breakdown stored for the spend logs carried no cache or reasoning lines for it. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CompletionTokensDetailsWrapper, CostPerToken + + logging_obj = Logging( + model="openai/onprem-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="custom-pricing-breakdown", + function_id="f", + ) + response = ModelResponse( + model="openai/onprem-model", + usage=Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ), + ) + + total = completion_cost( + completion_response=response, + model="openai/onprem-model", + custom_llm_provider="openai", + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ), + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(800 * 1e-7) + assert logging_obj.cost_breakdown["cache_creation_cost"] == pytest.approx(100 * 1.25e-6) + assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(200 * 2e-6) + assert total == pytest.approx(100 * 1e-6 + 800 * 1e-7 + 100 * 1.25e-6 + 500 * 2e-6) + + def test_cost_per_token_per_second_pricing(monkeypatch): """ Models priced by duration (input/output_cost_per_second) with no per-token rates diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6fb08445aff..6a9c86cc1b5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3301,11 +3301,14 @@ export interface paths { * - model: Model name (e.g., "gpt-4", "claude-3-opus") * - input_tokens: Expected input tokens per request * - output_tokens: Expected output tokens per request + * - cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional) + * - cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional) + * - reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional) * - num_requests_per_day: Number of requests per day (optional) * - num_requests_per_month: Number of requests per month (optional) * * Returns cost breakdown including: - * - Per-request costs (input, output, margin) + * - Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares) * - Daily costs (if num_requests_per_day provided) * - Monthly costs (if num_requests_per_month provided) * @@ -3314,7 +3317,9 @@ export interface paths { * { * "model": "gpt-4", * "input_tokens": 1000, + * "cache_read_input_tokens": 800, * "output_tokens": 500, + * "reasoning_tokens": 200, * "num_requests_per_day": 100, * "num_requests_per_month": 3000 * } @@ -26392,6 +26397,18 @@ export interface components { * @description Request body for /cost/estimate endpoint. */ CostEstimateRequest: { + /** + * Cache Creation Input Tokens + * @description Input tokens written to the prompt cache; counted within input_tokens + * @default 0 + */ + cache_creation_input_tokens: number; + /** + * Cache Read Input Tokens + * @description Input tokens read from the prompt cache; counted within input_tokens + * @default 0 + */ + cache_read_input_tokens: number; /** * Input Tokens * @description Expected input tokens per request @@ -26417,17 +26434,65 @@ export interface components { * @description Expected output tokens per request */ output_tokens: number; + /** + * Reasoning Tokens + * @description Reasoning tokens the model emits; counted within output_tokens + * @default 0 + */ + reasoning_tokens: number; }; /** * CostEstimateResponse * @description Response body for /cost/estimate endpoint. */ CostEstimateResponse: { + /** + * Cache Creation Cost Per Request + * @description Cache-write share of input_cost_per_request + * @default 0 + */ + cache_creation_cost_per_request: number; + /** + * Cache Creation Input Token Cost + * @description Rate billed per cache-write token + */ + cache_creation_input_token_cost?: number | null; + /** + * Cache Creation Input Tokens + * @default 0 + */ + cache_creation_input_tokens: number; + /** + * Cache Read Cost Per Request + * @description Cache-read share of input_cost_per_request + * @default 0 + */ + cache_read_cost_per_request: number; + /** + * Cache Read Input Token Cost + * @description Rate billed per cache-read token + */ + cache_read_input_token_cost?: number | null; + /** + * Cache Read Input Tokens + * @default 0 + */ + cache_read_input_tokens: number; /** * Cost Per Request * @description Total cost per request (includes margin) */ cost_per_request: number; + /** + * Daily Cache Creation Cost + * @description Cache-write share of daily_input_cost + */ + daily_cache_creation_cost?: number | null; + /** + * Daily Cache Read Cost + * @description Cache-read share of daily_input_cost + */ + daily_cache_read_cost?: number | null; /** * Daily Cost * @description Total daily cost (includes margin) @@ -26448,6 +26513,11 @@ export interface components { * @description Daily output token cost */ daily_output_cost?: number | null; + /** + * Daily Reasoning Cost + * @description Reasoning share of daily_output_cost + */ + daily_reasoning_cost?: number | null; /** * Input Cost Per Request * @description Input token cost per request (before margin) @@ -26465,6 +26535,16 @@ export interface components { margin_cost_per_request: number; /** Model */ model: string; + /** + * Monthly Cache Creation Cost + * @description Cache-write share of monthly_input_cost + */ + monthly_cache_creation_cost?: number | null; + /** + * Monthly Cache Read Cost + * @description Cache-read share of monthly_input_cost + */ + monthly_cache_read_cost?: number | null; /** * Monthly Cost * @description Total monthly cost (includes margin) @@ -26485,10 +26565,20 @@ export interface components { * @description Monthly output token cost */ monthly_output_cost?: number | null; + /** + * Monthly Reasoning Cost + * @description Reasoning share of monthly_output_cost + */ + monthly_reasoning_cost?: number | null; /** Num Requests Per Day */ num_requests_per_day?: number | null; /** Num Requests Per Month */ num_requests_per_month?: number | null; + /** + * Output Cost Per Reasoning Token + * @description Rate billed per reasoning token + */ + output_cost_per_reasoning_token?: number | null; /** * Output Cost Per Request * @description Output token cost per request (before margin) @@ -26500,6 +26590,17 @@ export interface components { output_tokens: number; /** Provider */ provider?: string | null; + /** + * Reasoning Cost Per Request + * @description Reasoning share of output_cost_per_request + * @default 0 + */ + reasoning_cost_per_request: number; + /** + * Reasoning Tokens + * @default 0 + */ + reasoning_tokens: number; }; /** CreateCredentialItem */ CreateCredentialItem: { From 7bff9bf9a2594d3cd5dd6a80a44e4e3b1cd38f61 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:22:32 -0700 Subject: [PATCH 015/157] fix(batches): handle provider cancellation and file cleanup gaps --- litellm/llms/bedrock/files/transformation.py | 70 ++++++++------ tests/e2e/batches/COVERAGE.md | 7 +- tests/e2e/batches/batch_cleanup.py | 63 +++++++++--- tests/e2e/batches/test_batch_cleanup.py | 69 +++++++++++++- tests/e2e/batches/test_batches_e2e.py | 14 +-- .../test_bedrock_files_transformation.py | 95 ++++++++++++++++--- 6 files changed, 258 insertions(+), 60 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 33b27943ad8..90b539ff37c 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -7,7 +7,7 @@ from contextlib import suppress from functools import cache from itertools import chain from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict +from typing import Any, Final, Literal, TypeAlias, TypedDict from urllib.parse import unquote import httpx @@ -60,11 +60,8 @@ from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id -# litellm_params key used to hand the SigV4-signed GET headers from -# `transform_file_content_request` to `validate_environment` (the only hook -# the shared file-content HTTP handler exposes for setting request headers). -# Same pattern as the `upload_url` handoff in `transform_create_file_request`. -S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" +S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" +S3_DELETE_FILE_ID_PARAM: Final = "_s3_delete_file_id" # litellm_params key carrying the size of the body uploaded to S3, handed from # `transform_create_file_request` to `transform_create_file_response`. @@ -291,7 +288,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) -> dict: result: Final[dict[str, object]] = {} result.update(headers) - signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None) + signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None) if isinstance(signed_headers, Mapping): result.update(signed_headers) # any-ok: untyped handoff headers # otherwise no extra headers - AWS credentials are handled by BaseAWSLLM @@ -1187,18 +1184,31 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + request: Final = self._transform_s3_file_request( + file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params + ) + litellm_params[S3_DELETE_FILE_ID_PARAM] = file_id + return request def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + if raw_response.status_code != 204: + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}", + headers=raw_response.headers, + ) + file_id: Final = litellm_params.get(S3_DELETE_FILE_ID_PARAM) + if not isinstance(file_id, str) or not file_id: + raise ValueError("Missing file id for Bedrock file deletion response") + return FileDeleted(id=file_id, deleted=True, object="file") def transform_list_files_request( self, @@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if not file_id: raise ValueError("file_id is required for Bedrock file content retrieval") + return self._transform_s3_file_request( + file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params + ) + + def _transform_s3_file_request( + self, + *, + file_id: str, + method: Literal["GET", "DELETE"], + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: s3_uri: Final = extract_s3_uri_from_file_id(file_id) bucket_name, object_key = _validate_file_id_against_configured_buckets( s3_uri=s3_uri, @@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) - # The shared file-content handler passes optional_params={}, so AWS - # credentials/region arrive via litellm_params here (unlike the upload - # path). s3_region_name wins over aws_region_name, same priority as - # get_complete_file_url above. - merged_params: Final[dict[str, object]] = {} - merged_params.update(litellm_params) - merged_params.update(optional_params) - request_params: Final = _BedrockS3RequestParams.model_validate(merged_params) + request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params}) region_preference: Final = request_params.s3_region_name or request_params.aws_region_name region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = ( + s3_endpoint_url: Final = ( request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" - litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body( api_base=url, aws_region_name=aws_region_name, request_params=request_params, + method=method, ) return url, {} - def _sign_s3_get_request( + def _sign_s3_request_without_body( self, api_base: str, aws_region_name: str, request_params: _BedrockS3RequestParams, + method: Literal["GET", "DELETE"] = "GET", ) -> dict[str, str]: - """ - SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT). - """ try: import hashlib @@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): empty_body_hash: Final = hashlib.sha256(b"").hexdigest() aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped - method="GET", + method=method, url=api_base, headers={"x-amz-content-sha256": empty_body_hash}, ) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index f95ea1f2649..ca44fc95e25 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -131,7 +131,12 @@ failures up to three times. Teardown attempts every registered cleanup before reporting failures as test errors. Already deleted files and batches that are terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes before input deletion: the ten-minute provider window plus a propagation margin. -Raw and model-encoded inputs can be deleted after cancellation is accepted +Accepted cancellation may still report validating or in_progress while the provider +updates its state. Raw and model-encoded batches are polled until cancelling or +terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes +output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE +restricted to the configured storage buckets and managed file prefixes. The low-RPM +test submits with its restricted key and cleans up with the test administrator key Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 722df0c29bc..9284882ad82 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -1,16 +1,17 @@ +from builtins import ExceptionGroup from collections.abc import Callable from itertools import count from time import monotonic, sleep from typing import Final, Protocol -from pydantic import BaseModel - from batch_client import BatchObject, FileDeleteResponse from capabilities import is_managed_id from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from pydantic import BaseModel CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) +BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"}) BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 BATCH_CANCEL_POLL_SECONDS: Final = 10.0 @@ -63,6 +64,7 @@ def cleanup_batch( *, key: str, provider: str | None = None, + delete_output_files: bool = False, wait: Callable[[float], None] = sleep, clock: Callable[[], float] = monotonic, ) -> None: @@ -72,18 +74,28 @@ def cleanup_batch( f"Retrieve batch {batch_id} for cleanup", ) if fetched.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, fetched, key=key, provider=provider) return if fetched.status == "cancelling" and not needs_terminal_state: return - if fetched.status != "cancelling": - result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) - if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): - cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") - assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( - f"Cancel batch {batch_id} left status {cancelled.status}" - ) - if not needs_terminal_state: - return + result: Final = ( + Success(status_code=200, data=fetched) + if fetched.status == "cancelling" + else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) + ) + conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409} + if not conflicted: + cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") + assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, ( + f"Cancel batch {batch_id} left status {cancelled.status}" + ) + if cancelled.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, cancelled, key=key, provider=provider) + return + if cancelled.status == "cancelling" and not needs_terminal_state: + return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS for current in ( _require_cleanup_success( @@ -93,11 +105,36 @@ def cleanup_batch( for _ in count() ): if current.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, current, key=key, provider=provider) return - assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" - if not needs_terminal_state: + assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), ( + f"Cancel batch {batch_id} left status {current.status}" + ) + if current.status == "cancelling" and not needs_terminal_state: return assert clock() < deadline, ( f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" ) wait(BATCH_CANCEL_POLL_SECONDS) + + +def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None: + errors: Final = tuple( + error + for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id)) + if file_id is not None and file_id != batch.input_file_id + if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None + ) + if errors: + raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors) + + +def _output_cleanup_error( + client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None +) -> Exception | None: + try: + cleanup_file(client, file_id, key=key, provider=provider) + except Exception as error: + return error + return None diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 66b7079ccc2..a875aee719b 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -4,7 +4,6 @@ from dataclasses import dataclass, field from typing import Final import pytest - from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form from capabilities import CAPABILITIES, Capability @@ -219,6 +218,74 @@ class TestBatchCancellation: cleanup_batch(client, "batch-1", key="test-key", provider="azure") client.calls.assert_done() + @pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID]) + @pytest.mark.parametrize("pending_status", ["validating", "in_progress"]) + def test_accepted_cancellation_waits_through_stale_provider_status( + self, batch_id: str, pending_status: str + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + iter( + ( + f"retrieve vertex_ai {batch_id}", + f"cancel vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + "delete vertex_ai file-1", + "delete key test-key", + ) + ) + ), + batches=iter((batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled"))), + cancellations=iter((batch(pending_status),)), + files=iter((deleted_file(),)), + ) + delays: Final = ExpectedCalls(iter((10.0, 10.0))) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai")) + manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays)) + manager.teardown() + client.calls.assert_done() + delays.assert_done() + + @pytest.mark.parametrize("output_delete_fails", [False, True]) + def test_batch_that_completed_before_cleanup_deletes_output_and_error_files( + self, output_delete_fails: bool + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + iter(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")) + ), + batches=iter( + ( + Success( + status_code=200, + data=BatchObject( + id="batch-1", + status="completed", + input_file_id="file-input", + output_file_id="file-output", + error_file_id="file-error", + ), + ), + ) + ), + files=iter( + ( + UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), + deleted_file(), + ) + ), + ) + if output_delete_fails: + with pytest.raises(ExceptionGroup, match="output cleanup failed"): + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + else: + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + client.calls.assert_done() + @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1b4a6ed266f..c4b699190b8 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -25,7 +25,7 @@ from datetime import datetime, timedelta, timezone import pytest from pydantic import BaseModel -from e2e_config import PROXY_BASE_URL, unique_marker +from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker from batch_cleanup import cleanup_batch, cleanup_file from batch_client import ( @@ -257,7 +257,9 @@ def test_batch_lifecycle( require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( - lambda: cleanup_batch(client, batch.id, key=key, provider=provider) + lambda: cleanup_batch( + client, batch.id, key=key, provider=provider, delete_output_files=cap.provider in {"openai", "azure"} + ) ) assert batch.id, f"create returned no batch id (body={created.body[:200]})" @@ -801,7 +803,7 @@ class TestBatchEnqueuedTokenLimit: """ def _upload_batch_file( - self, client: BatchClient, resources: ResourceManager, key: str + self, client: BatchClient, resources: ResourceManager, key: str, *, cleanup_key: str | None = None ) -> FileObject: file = unwrap( client.upload_file( @@ -811,7 +813,7 @@ class TestBatchEnqueuedTokenLimit: key=key, ) ) - resources.defer(lambda: cleanup_file(client, file.id, key=key)) + resources.defer(lambda: cleanup_file(client, file.id, key=cleanup_key or key)) return file def _generate_enqueued_key( @@ -848,7 +850,7 @@ class TestBatchEnqueuedTokenLimit: marker="rpm", rpm_limit=BATCH_RL_RPM_LIMIT, ) - file = self._upload_batch_file(client, resources, key) + file = self._upload_batch_file(client, resources, key, cleanup_key=MASTER_KEY) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -859,7 +861,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) + resources.defer(lambda: cleanup_batch(client, batch.id, key=MASTER_KEY, delete_output_files=True)) @pytest.mark.covers( "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 541c0db15d8..2c02a58663e 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -5,6 +5,8 @@ Test bedrock files transformation functionality import json import os from collections.abc import Mapping +from contextlib import AsyncExitStack, closing +from typing import Final from unittest.mock import MagicMock from urllib.parse import unquote, urlparse @@ -1855,6 +1857,77 @@ class TestBedrockBatchNonChatEndpointRecords: ] +class TestBedrockFileDeletion: + S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl" + URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl" + + def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with respx.mock, closing(HTTPHandler()) as client: + route: Final = respx.delete(self.URL).mock(return_value=httpx.Response(204)) + deleted: Final = litellm.file_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + request: Final = route.calls[0].request + assert request.content == b"" + signed: Final = AWSRequest(method="DELETE", url=self.URL, headers={ + "X-Amz-Date": request.headers["X-Amz-Date"], + "X-Amz-Content-SHA256": request.headers["X-Amz-Content-SHA256"], + }) + signed.context["timestamp"] = request.headers["X-Amz-Date"] + auth: Final = S3SigV4Auth(Credentials("AKIAEXAMPLE", "test-secret"), "s3", "us-west-2") + signature: Final = auth.signature(auth.string_to_sign(signed, auth.canonical_request(signed)), signed) + assert request.headers["Authorization"].endswith(f"Signature={signature}") + assert deleted.id == self.S3_URI and deleted.deleted is True + + @pytest.mark.asyncio + async def test_adelete_file_propagates_s3_errors(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + async with AsyncExitStack() as stack: + client: Final = AsyncHTTPHandler() + stack.push_async_callback(client.close) + with respx.mock: + route: Final = respx.delete(self.URL).mock( + return_value=httpx.Response(403, content=b"AccessDenied") + ) + from litellm.llms.bedrock.common_utils import BedrockError + + with pytest.raises(BedrockError, match="AccessDenied"): + await litellm.afile_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + + @pytest.mark.parametrize("file_id, message", [ + ("s3://other-bucket/litellm-bedrock-files-model-abc.jsonl", "configured storage bucket"), + ("s3://my-bucket/private/data.jsonl", "LiteLLM-managed"), + ]) + def test_delete_rejects_untrusted_objects_before_signing( + self, file_id: str, message: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with pytest.raises(ValueError, match=message): + BedrockFilesConfig().transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params={}) + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" @@ -1873,7 +1946,7 @@ class TestBedrockFileContentTransformation: import hashlib from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -1889,7 +1962,7 @@ class TestBedrockFileContentTransformation: assert url == self.EXPECTED_URL assert params == {} - signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM] + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] content_hashes = { value for name, value in signed_headers.items() @@ -2139,7 +2212,7 @@ class TestBedrockFileContentTransformation: def test_s3_region_name_wins_for_content_signing(self, monkeypatch): """s3_region_name must override aws_region_name for both the URL and the signature.""" from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2154,17 +2227,17 @@ class TestBedrockFileContentTransformation: ) assert url.startswith("https://s3.eu-west-1.amazonaws.com/") - authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"] + authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"] assert "/eu-west-1/s3/aws4_request" in authorization def test_validate_environment_merges_and_pops_signed_get_headers(self): from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) litellm_params = { - S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} + S3_SIGNED_REQUEST_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} } headers = BedrockFilesConfig().validate_environment( @@ -2179,7 +2252,7 @@ class TestBedrockFileContentTransformation: "x-custom": "kept", "Authorization": "AWS4-HMAC-SHA256 test", } - assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params + assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params def test_transform_file_content_response_wraps_binary_content(self): import httpx @@ -2379,7 +2452,7 @@ class TestBedrockFilesS3SignatureEncoding: self, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2402,7 +2475,7 @@ class TestBedrockFilesS3SignatureEncoding: method="GET", url=url, body=None, - headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM], + headers=litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM], ) @@ -2457,7 +2530,7 @@ def test_sign_s3_request_assumes_role_with_external_id(monkeypatch): assert "ASIAFILESPUTROLE" in authorization -def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): +def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch): """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request.""" import datetime from unittest.mock import patch @@ -2504,7 +2577,7 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): assert request_params.aws_external_id == "external-id-files-get" with patch.object(boto3, "client", return_value=FakeSTSClient()): - signed_headers = BedrockFilesConfig()._sign_s3_get_request( + signed_headers = BedrockFilesConfig()._sign_s3_request_without_body( api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", aws_region_name="us-east-1", request_params=request_params, From 4ab5719ff9e0770ecb9f2d1b53c4caf58f19e5db Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:36:27 -0700 Subject: [PATCH 016/157] test(batches): use immutable expectations with explicit test doubles --- litellm/files/main.py | 2 +- tests/e2e/batches/test_batch_cleanup.py | 187 ++++++++++++------------ 2 files changed, 93 insertions(+), 96 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 19da77b7364..218518eb3cd 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -31,7 +31,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] import litellm from litellm import get_secret_str diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index a875aee719b..d0038139dcf 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -1,7 +1,7 @@ from builtins import ExceptionGroup -from collections.abc import Iterator -from dataclasses import dataclass, field +from collections.abc import Callable from typing import Final +from unittest.mock import Mock, call import pytest from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result @@ -15,35 +15,43 @@ MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" -@dataclass(frozen=True, slots=True) class ExpectedCalls[T]: - values: Iterator[T] + def __init__(self, values: tuple[T, ...]) -> None: + self.values: Final = values + self.recorder: Final = Mock() def __call__(self, value: T) -> None: - assert next(self.values, None) == value + self.recorder(value) def assert_done(self) -> None: - assert tuple(self.values) == () + assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values) -@dataclass(frozen=True, slots=True) class CleanupClient: - calls: ExpectedCalls[str] - files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) - batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) - cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + def __init__( + self, + *, + calls: ExpectedCalls[str], + files: tuple[Result[FileDeleteResponse], ...] = (), + batches: tuple[Result[BatchObject], ...] = (), + cancellations: tuple[Result[BatchObject], ...] = (), + ) -> None: + self.calls: Final = calls + self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files) + self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches) + self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations) def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: self.calls(f"delete {provider} {file_id}") - return next(self.files) + return self.file_response() def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: self.calls(f"retrieve {provider} {batch_id}") - return next(self.batches) + return self.batch_response() def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: self.calls(f"cancel {provider} {batch_id}") - return next(self.cancellations) + return self.cancel_response() def generate_key(self, body: KeyGenerateBody) -> str: return "test-key" @@ -68,17 +76,15 @@ class TestFileCleanup: response: Final = Success( status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) ) - client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete None {MANAGED_FILE_ID}",))), files=iter((response,)) - ) + client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,)) cleanup_file(client, MANAGED_FILE_ID, key="test-key") client.calls.assert_done() @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete None {file_id}",))), - files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),)), + calls=ExpectedCalls((f"delete None {file_id}",)), + files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),), ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, file_id, key="test-key") @@ -88,15 +94,15 @@ class TestFileCleanup: def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete {expected_provider} file-1",))), files=iter((deleted_file(),)) + calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),) ) cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) client.calls.assert_done() def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete azure file-1", "delete key test-key"))), - files=iter((UnknownApiError(status_code=403, body="secret response"),)), + calls=ExpectedCalls(("delete azure file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="secret response"),), ) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() @@ -109,7 +115,7 @@ class TestFileCleanup: def test_success_response_must_confirm_deletion(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete None file-1",))), files=iter((deleted_file(deleted=False),)) + calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),) ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, "file-1", key="test-key") @@ -117,16 +123,16 @@ class TestFileCleanup: def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete azure file-1",))), - files=iter((UnknownApiError(status_code=404, body="missing"),)), + calls=ExpectedCalls(("delete azure file-1",)), + files=(UnknownApiError(status_code=404, body="missing"),), ) cleanup_file(client, "file-1", key="test-key", provider="azure") client.calls.assert_done() def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete None file-1", "delete key test-key"))), - files=iter((UnknownApiError(status_code=403, body="forbidden"),)), + calls=ExpectedCalls(("delete None file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="forbidden"),), ) manager: Final = ResourceManager(client=client) key: Final = manager.key() @@ -141,37 +147,39 @@ class TestCleanupRetries: [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")], ) def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: - outcomes: Final = iter((failure, deleted_file())) - delays: Final = ExpectedCalls(iter((1.0,))) - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls((1.0,)) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) assert isinstance(result, Success) and result.data.deleted delays.assert_done() def test_persistent_error_has_bounded_retries(self) -> None: failure: Final = UnknownApiError(status_code=503, body="unavailable") - outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) - delays: Final = ExpectedCalls(iter(CLEANUP_DELAYS)) - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) + outcomes: Final = Mock(return_value=failure) + delays: Final = ExpectedCalls(CLEANUP_DELAYS) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) assert result is failure delays.assert_done() - assert next(outcomes, None) is None + assert outcomes.call_count == len(CLEANUP_DELAYS) + 1 def test_permanent_error_is_not_retried(self) -> None: failure: Final = UnknownApiError(status_code=403, body="forbidden") - outcomes: Final = iter((failure, deleted_file())) - delays: Final = ExpectedCalls[float](iter(())) - assert cleanup_result(lambda: next(outcomes), wait=delays) is failure + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls[float](()) + assert cleanup_result(outcomes, wait=delays) is failure delays.assert_done() - assert isinstance(next(outcomes), Success) + assert outcomes.call_count == 1 class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"retrieve None {MANAGED_BATCH_ID}",) * 3)), - batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled"))), + calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3), + batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")), ) - delays: Final = ExpectedCalls(iter((10.0,))) + delays: Final = ExpectedCalls((10.0,)) cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays) client.calls.assert_done() delays.assert_done() @@ -179,23 +187,22 @@ class TestBatchCancellation: def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: client: Final = CleanupClient( calls=ExpectedCalls( - iter( - ( - f"retrieve None {MANAGED_BATCH_ID}", - f"retrieve None {MANAGED_BATCH_ID}", - "delete None file-1", - "delete key test-key", - ) + ( + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", + "delete None file-1", + "delete key test-key", ) ), - batches=iter((batch("cancelling"), batch("cancelling"))), - files=iter((deleted_file(),)), + batches=(batch("cancelling"), batch("cancelling")), + files=(deleted_file(),), ) - ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) + times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS) + ticks: Final[Callable[[], float]] = Mock(side_effect=times) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) - manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=lambda: next(ticks))) + manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks)) with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) @@ -203,17 +210,15 @@ class TestBatchCancellation: @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: - client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve None batch-1",))), batches=iter((batch(status),)) - ) + client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),)) cleanup_batch(client, "batch-1", key="test-key") client.calls.assert_done() def test_active_batch_is_cancelled_through_its_provider(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve azure batch-1", "cancel azure batch-1"))), - batches=iter((batch("in_progress"), batch("cancelled"))), - cancellations=iter((batch("cancelling"),)), + calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")), + batches=(batch("in_progress"), batch("cancelled")), + cancellations=(batch("cancelling"),), ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") client.calls.assert_done() @@ -225,23 +230,21 @@ class TestBatchCancellation: ) -> None: client: Final = CleanupClient( calls=ExpectedCalls( - iter( - ( - f"retrieve vertex_ai {batch_id}", - f"cancel vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - "delete vertex_ai file-1", - "delete key test-key", - ) + ( + f"retrieve vertex_ai {batch_id}", + f"cancel vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + "delete vertex_ai file-1", + "delete key test-key", ) ), - batches=iter((batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled"))), - cancellations=iter((batch(pending_status),)), - files=iter((deleted_file(),)), + batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")), + cancellations=(batch(pending_status),), + files=(deleted_file(),), ) - delays: Final = ExpectedCalls(iter((10.0, 10.0))) + delays: Final = ExpectedCalls((10.0, 10.0)) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai")) @@ -255,28 +258,22 @@ class TestBatchCancellation: self, output_delete_fails: bool ) -> None: client: Final = CleanupClient( - calls=ExpectedCalls( - iter(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")) - ), - batches=iter( - ( - Success( - status_code=200, - data=BatchObject( - id="batch-1", - status="completed", - input_file_id="file-input", - output_file_id="file-output", - error_file_id="file-error", - ), + calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")), + batches=( + Success( + status_code=200, + data=BatchObject( + id="batch-1", + status="completed", + input_file_id="file-input", + output_file_id="file-output", + error_file_id="file-error", ), - ) + ), ), - files=iter( - ( - UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), - deleted_file(), - ) + files=( + UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), + deleted_file(), ), ) if output_delete_fails: @@ -289,9 +286,9 @@ class TestBatchCancellation: @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"))), - batches=iter((batch("in_progress"), batch(status))), - cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), + calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")), + batches=(batch("in_progress"), batch(status)), + cancellations=(UnknownApiError(status_code=409, body="conflict"),), ) if status == "completed": cleanup_batch(client, "batch-1", key="test-key") From 43a02e2dbcd485c3cc5dd31d8b11161da650c356 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 7 Sep 2026 16:53:41 -0700 Subject: [PATCH 017/157] fix(proxy): report the billed token rates in /cost/estimate The rate fields reported base cost-map prices while the cost lines were billed at the token tier, off-peak window and regional multipliers the calculator picks for the request, so a line did not always equal tokens times its reported rate. get_billed_token_rates now resolves the rates once, the token-type breakdown and the endpoint both read from it, and a tiered-model test asserts every line equals its token count times the rate reported next to it Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- .../litellm_core_utils/llm_cost_calc/utils.py | 206 ++++++++++++------ litellm/proxy/_types.py | 6 +- .../cost_tracking_settings.py | 63 ++---- .../llm_cost_calc/test_llm_cost_calc_utils.py | 49 +++++ .../test_cost_tracking_settings.py | 52 ++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 +- 6 files changed, 258 insertions(+), 128 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8d53c8da3e6..2bb15fb4c48 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1329,16 +1329,118 @@ def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetai ) -def _custom_pricing_token_type_breakdown(usage: Usage, custom_cost_per_token: CostPerToken) -> TokenTypeCostBreakdown: +@dataclass(frozen=True, slots=True) +class BilledTokenRates: + """Per-token rates one request's usage bills at, after token tiers, off-peak windows and the + regional multipliers the totals apply, so each cost line equals its token count times its rate.""" + + input_cost_per_token: float + output_cost_per_token: float + cache_read_input_token_cost: float + cache_creation_input_token_cost: float + cache_creation_input_token_cost_above_1hr: float + output_cost_per_reasoning_token: float + + def scaled(self, multiplier: float) -> "BilledTokenRates": + if multiplier == 1.0: + return self + return BilledTokenRates( + input_cost_per_token=self.input_cost_per_token * multiplier, + output_cost_per_token=self.output_cost_per_token * multiplier, + cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, + cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, + cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, + output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, + ) + + +def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates: """Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" input_rate: Final = custom_cost_per_token["input_cost_per_token"] - cache_read_tokens, cache_creation_tokens, _ = _cache_token_counts(usage) - return TokenTypeCostBreakdown( - reasoning_cost=float(_reasoning_token_count(usage)) * custom_cost_per_token["output_cost_per_token"], - cache_read_cost=float(cache_read_tokens) * custom_cost_per_token.get("cache_read_input_token_cost", input_rate), - cache_creation_cost=float(cache_creation_tokens) - * custom_cost_per_token.get("cache_creation_input_token_cost", input_rate), + output_rate: Final = custom_cost_per_token["output_cost_per_token"] + cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate) + return BilledTokenRates( + input_cost_per_token=input_rate, + output_cost_per_token=output_rate, + cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_creation_input_token_cost=cache_creation_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_rate, + output_cost_per_reasoning_token=output_rate, + ) + + +def _cost_map_billed_rates( + model_info: ModelInfo, + usage: Usage, + custom_llm_provider: str | None, + service_tier: str | None, + data_residency: str | None, + vertex_location: str | None, + current_time: datetime | None, +) -> BilledTokenRates: + billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + ( + prompt_base_cost, + completion_base_cost, + cache_creation_cost_rate, + cache_creation_cost_above_1hr_rate, + cache_read_cost_rate, + ) = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier=service_tier, + current_time=billing_time, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), + ) + reasoning_rate: Final = _resolve_billed_reasoning_rate( + model_info=model_info, + usage=usage, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + current_time=billing_time, + ) + multiplier: Final = ( + _get_regional_uplift_multiplier(model_info, data_residency) + * get_vertex_regional_endpoint_uplift(model_info, vertex_location) + * get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + ) + return BilledTokenRates( + input_cost_per_token=prompt_base_cost, + output_cost_per_token=completion_base_cost, + cache_read_input_token_cost=cache_read_cost_rate, + cache_creation_input_token_cost=cache_creation_cost_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, + output_cost_per_reasoning_token=reasoning_rate, + ).scaled(multiplier) + + +def get_billed_token_rates( + model: str, + custom_llm_provider: str | None, + usage: Usage, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, +) -> BilledTokenRates | None: + """Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type + breakdown resolve them. None when the model's pricing cannot be resolved.""" + if custom_cost_per_token is not None: + return _custom_pricing_rates(custom_cost_per_token) + try: + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + return _cost_map_billed_rates( + model_info=model_info, + usage=usage, + custom_llm_provider=custom_llm_provider, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, ) @@ -1360,77 +1462,39 @@ def get_token_type_cost_breakdown( cost calculators bypass ``generic_cost_per_token``, because cache tokens always land on ``prompt_tokens_details`` (via the Usage constructor and provider transformations) and reasoning tokens on ``completion_tokens_details``. It reuses - the same rate-resolution primitives as the total-cost path so the breakdown can - never drift from the totals. A deployment billed by ``custom_cost_per_token`` is - priced from those flat rates instead of the cost map, for the same reason. + the same rate resolution as the total-cost path (``get_billed_token_rates``) so the + breakdown can never drift from the totals. A deployment billed by + ``custom_cost_per_token`` is priced from those flat rates instead of the cost map and, + like its totals, bills cache writes flat rather than by their 5m/1h split. Returns zeros (never raises) when the model or its pricing cannot be resolved. """ - if custom_cost_per_token is not None: - return _custom_pricing_token_type_breakdown(usage=usage, custom_cost_per_token=custom_cost_per_token) - - try: - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: + rates: Final = get_billed_token_rates( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, + custom_cost_per_token=custom_cost_per_token, + ) + if rates is None: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) - ( - _prompt_base_cost, - completion_base_cost, - cache_creation_cost_rate, - cache_creation_cost_above_1hr_rate, - cache_read_cost_rate, - ) = _get_token_base_cost( - model_info=model_info, - usage=usage, - service_tier=service_tier, - current_time=billing_time, - threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), - ) - - reasoning_rate: Final = _resolve_billed_reasoning_rate( - model_info=model_info, - usage=usage, - service_tier=service_tier, - completion_base_cost=completion_base_cost, - current_time=billing_time, - ) - reasoning_cost = float(_reasoning_token_count(usage)) * reasoning_rate - cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) - cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate - cache_creation_cost = calculate_cache_writing_cost( - cache_creation_tokens=cache_creation_tokens, - cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr_rate, - cache_creation_cost=cache_creation_cost_rate, + cache_creation_cost: Final = ( + float(cache_creation_tokens) * rates.cache_creation_input_token_cost + if custom_cost_per_token is not None + else calculate_cache_writing_cost( + cache_creation_tokens=cache_creation_tokens, + cache_creation_token_details=cache_creation_token_details, + cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr, + cache_creation_cost=rates.cache_creation_input_token_cost, + ) ) - - # Apply the same flat regional-processing uplift the totals get, so per-type - # costs stay reconciled with input_cost/output_cost for regionalized OpenAI hosts. - uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) - if uplift != 1.0: - reasoning_cost *= uplift - cache_read_cost *= uplift - cache_creation_cost *= uplift - - vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) - if vertex_uplift != 1.0: - reasoning_cost *= vertex_uplift - cache_read_cost *= vertex_uplift - cache_creation_cost *= vertex_uplift - - # Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals - # apply, so cache and reasoning line items stay reconciled with them. - geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) - if geo_multiplier != 1.0: - reasoning_cost *= geo_multiplier - cache_read_cost *= geo_multiplier - cache_creation_cost *= geo_multiplier - return TokenTypeCostBreakdown( - reasoning_cost=reasoning_cost, - cache_read_cost=cache_read_cost, + reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, + cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, cache_creation_cost=cache_creation_cost, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 65b7d409d7c..3c1cd9bfc69 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5197,9 +5197,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): default=None, description="Cache-write share of monthly_input_cost" ) monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost") - # Pricing info - input_cost_per_token: float | None = None - output_cost_per_token: float | None = None + # Pricing info: the rates this request's usage bills at, after token tiers and regional multipliers + input_cost_per_token: float | None = Field(default=None, description="Rate billed per input token") + output_cost_per_token: float | None = Field(default=None, description="Rate billed per output token") cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token") cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token") output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token") diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 51f31a757cd..17d82fd17e3 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -20,6 +20,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.llm_cost_calc.utils import get_billed_token_rates from litellm.proxy._types import ( CommonProxyErrors, CostEstimateRequest, @@ -170,44 +171,6 @@ def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) - ) -@dataclass(frozen=True, slots=True) -class EffectiveTokenRates: - input_cost_per_token: float | None - output_cost_per_token: float | None - cache_read_input_token_cost: float | None - cache_creation_input_token_cost: float | None - output_cost_per_reasoning_token: float | None - - -def _custom_token_rates(custom_cost_per_token: CostPerToken) -> EffectiveTokenRates: - input_rate: Final = custom_cost_per_token["input_cost_per_token"] - output_rate: Final = custom_cost_per_token["output_cost_per_token"] - return EffectiveTokenRates( - input_cost_per_token=input_rate, - output_cost_per_token=output_rate, - cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), - cache_creation_input_token_cost=custom_cost_per_token.get("cache_creation_input_token_cost", input_rate), - output_cost_per_reasoning_token=output_rate, - ) - - -def _cost_map_token_rates(model_info: ModelInfo | None) -> EffectiveTokenRates: - """Base rates the cost calculator bills flat usage at: a cost-map model without a cache price - bills cache tokens at zero, and one without a reasoning price bills reasoning at the output rate.""" - if model_info is None: - return EffectiveTokenRates(None, None, None, None, None) - sources: Final = (model_info,) - output_rate: Final = _configured_price("output_cost_per_token", sources) - reasoning_rate: Final = _configured_price("output_cost_per_reasoning_token", sources) - return EffectiveTokenRates( - input_cost_per_token=_configured_price("input_cost_per_token", sources), - output_cost_per_token=output_rate, - cache_read_input_token_cost=_configured_price("cache_read_input_token_cost", sources) or 0.0, - cache_creation_input_token_cost=_configured_price("cache_creation_input_token_cost", sources) or 0.0, - output_cost_per_reasoning_token=output_rate if reasoning_rate is None else reasoning_rate, - ) - - def _usage_for_estimate(request: CostEstimateRequest) -> Usage: cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens return Usage( @@ -654,7 +617,8 @@ async def estimate_cost( verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model) - mock_response: Final = ModelResponse(model=resolved_model, usage=_usage_for_estimate(request)) + usage: Final = _usage_for_estimate(request) + mock_response: Final = ModelResponse(model=resolved_model, usage=usage) # Create a logging object to capture cost breakdown litellm_logging_obj: Final = LiteLLMLoggingObj( @@ -688,12 +652,13 @@ async def estimate_cost( daily: Final = per_request.times(request.num_requests_per_day) monthly: Final = per_request.times(request.num_requests_per_month) - model_info: Final = _lookup_model_info(resolved_model) - rates: Final = ( - _custom_token_rates(resolved.custom_cost_per_token) - if resolved.custom_cost_per_token is not None - else _cost_map_token_rates(model_info) + rates: Final = get_billed_token_rates( + model=resolved_model, + custom_llm_provider=resolved_provider, + usage=usage, + custom_cost_per_token=resolved.custom_cost_per_token, ) + model_info: Final = _lookup_model_info(resolved_model) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider @@ -727,10 +692,10 @@ async def estimate_cost( monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None, monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None, monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None, - input_cost_per_token=rates.input_cost_per_token, - output_cost_per_token=rates.output_cost_per_token, - cache_read_input_token_cost=rates.cache_read_input_token_cost, - cache_creation_input_token_cost=rates.cache_creation_input_token_cost, - output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token, + input_cost_per_token=rates.input_cost_per_token if rates is not None else None, + output_cost_per_token=rates.output_cost_per_token if rates is not None else None, + cache_read_input_token_cost=rates.cache_read_input_token_cost if rates is not None else None, + cache_creation_input_token_cost=rates.cache_creation_input_token_cost if rates is not None else None, + output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token if rates is not None else None, provider=custom_llm_provider, ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7065003839b..4de3c059e63 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -27,6 +27,7 @@ from litellm.types.utils import ( ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, PromptTokensDetailsResult, TokenRates, @@ -38,6 +39,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( apply_off_peak_pricing, calculate_cache_writing_cost, generic_cost_per_token, + get_billed_token_rates, get_token_type_cost_breakdown, ) from litellm.types.utils import CacheCreationTokenDetails, Usage @@ -3938,6 +3940,53 @@ def test_token_type_cost_breakdown_reconciles_with_custom_pricing_totals(): assert 300 * 2e-6 + breakdown.reasoning_cost == pytest.approx(completion_cost) +def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "tiered-cache-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200_000, cache_creation_tokens=10_000), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + rates = get_billed_token_rates(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + breakdown = get_token_type_cost_breakdown(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + + assert rates == BilledTokenRates( + input_cost_per_token=6e-6, + output_cost_per_token=3e-5, + cache_read_input_token_cost=6e-7, + cache_creation_input_token_cost=7.5e-6, + cache_creation_input_token_cost_above_1hr=0.0, + output_cost_per_reasoning_token=3e-5, + ) + assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) + assert breakdown.cache_creation_cost == pytest.approx(10_000 * rates.cache_creation_input_token_cost) + assert breakdown.reasoning_cost == pytest.approx(200 * rates.output_cost_per_reasoning_token) + + +def test_billed_token_rates_are_none_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + assert get_billed_token_rates(model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage) is None + + def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 6d3ef2ffea9..9847c4092df 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -813,9 +813,7 @@ async def _estimate(mock_router: MagicMock | None, model: str = AN_ALIAS, **over request = CostEstimateRequest( model=model, - input_tokens=INPUT_TOKENS, - output_tokens=OUTPUT_TOKENS, - **overrides, + **{"input_tokens": INPUT_TOKENS, "output_tokens": OUTPUT_TOKENS, **overrides}, ) with patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point "litellm.proxy.proxy_server.llm_router", mock_router @@ -1062,6 +1060,54 @@ class TestEstimateCostCacheAndReasoningTokens: assert response.cache_read_input_token_cost == pytest.approx(5e-7) assert response.cache_creation_input_token_cost == pytest.approx(6.25e-6) + @pytest.mark.asyncio + async def test_a_tiered_model_reports_the_rates_its_lines_were_billed_at(self, monkeypatch): + """Above a token tier the calculator bills every line at the tier's rate, so the reported + rates must be the tier's too: each line equals its token count times the rate next to it.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=A_MAPPED_MODEL, + input_tokens=250_000, + cache_read_input_tokens=200_000, + cache_creation_input_tokens=10_000, + output_tokens=1_000, + reasoning_tokens=200, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_creation_input_token_cost == pytest.approx(7.5e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(3e-5) + assert response.cache_read_cost_per_request == pytest.approx(200_000 * response.cache_read_input_token_cost) + assert response.cache_creation_cost_per_request == pytest.approx( + 10_000 * response.cache_creation_input_token_cost + ) + assert response.reasoning_cost_per_request == pytest.approx(200 * response.output_cost_per_reasoning_token) + assert response.input_cost_per_request == pytest.approx( + 40_000 * response.input_cost_per_token + + response.cache_read_cost_per_request + + response.cache_creation_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + class TestCostEstimateRequestTokenSubsets: def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6a9c86cc1b5..46cb5ad892a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26523,7 +26523,10 @@ export interface components { * @description Input token cost per request (before margin) */ input_cost_per_request: number; - /** Input Cost Per Token */ + /** + * Input Cost Per Token + * @description Rate billed per input token + */ input_cost_per_token?: number | null; /** Input Tokens */ input_tokens: number; @@ -26584,7 +26587,10 @@ export interface components { * @description Output token cost per request (before margin) */ output_cost_per_request: number; - /** Output Cost Per Token */ + /** + * Output Cost Per Token + * @description Rate billed per output token + */ output_cost_per_token?: number | null; /** Output Tokens */ output_tokens: number; From 7d3b68fea5a6343781401e52f40b02f6c581541c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:59:53 -0700 Subject: [PATCH 018/157] fix(files): preserve managed deletion routing and response identity --- .../proxy/hooks/managed_files.py | 13 ++- tests/e2e/batches/COVERAGE.md | 3 + .../proxy/test_managed_files_hook.py | 104 ++++++++++++++++++ 3 files changed, 118 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc1eb6cebc2..6e0bb0da3f4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1779,7 +1779,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Remove conflicting keys from data to avoid duplicate keyword arguments filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): - delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + delete_data = { + **{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"}, + **( + {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} + if credentials is not None + else {} + ), + } + delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data) stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) @@ -1790,7 +1799,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): prom_logger.record_managed_file_deleted(result="success") if stored_file_object: - return stored_file_object + return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id}) elif delete_response: delete_response.id = file_id return delete_response diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca44fc95e25..919c39f21a2 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -138,6 +138,9 @@ output and error files returned by terminal batches. Bedrock deletion uses a sig restricted to the configured storage buckets and managed file prefixes. The low-RPM test submits with its restricted key and cleans up with the test administrator key +Managed deletion forwards the deployment's trusted bucket configuration and returns +the requested managed file ID even when stored output metadata carries a provider ID + Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a fallback for interrupted runs: immediate deletion remains the normal cleanup. diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 091b958d7c3..48fceb50403 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -1095,6 +1095,110 @@ async def test_afile_content_passes_trusted_model_credentials_to_router(): assert trusted_credentials["s3_bucket_name"] == "my-bucket" +def _managed_deletion_file_id(provider_file_id): + from litellm.types.utils import SpecialEnums + + value = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "test-file", "batch-model", provider_file_id, "model-123" + ) + return base64.urlsafe_b64encode(value.encode()).decode().rstrip("=") + + +def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object): + from litellm.caching import DualCache + from litellm.models.managed_files import LiteLLM_ManagedFileTable + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + row = LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + model_mappings={"model-123": provider_file_id}, + flat_model_file_ids=[provider_file_id], + file_object=file_object, + ) + table = MagicMock( + find_first=AsyncMock(return_value=row), + delete=AsyncMock(), + ) + return _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)), + ), table + + +@pytest.mark.asyncio +async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete(monkeypatch): + import httpx + import respx + + from litellm import Router + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + router = Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + "model_info": {"id": "model-123"}, + } + ], + num_retries=0, + ) + s3_uri = "s3://my-bucket/litellm-bedrock-files/input.jsonl" + unified_file_id = _managed_deletion_file_id(s3_uri) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, s3_uri, None) + with respx.mock: + route = respx.delete( + "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl" + ).mock(return_value=httpx.Response(204)) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert len(route.calls) == 1 + assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert response.id == unified_file_id + assert response.deleted is True + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + +@pytest.mark.asyncio +async def test_afile_delete_returns_managed_id_for_stored_provider_output(): + from openai.types import FileDeleted + + provider_file_id = "file-error-output" + unified_file_id = _managed_deletion_file_id(provider_file_id) + stored_file = _make_file_object(provider_file_id) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, provider_file_id, stored_file) + router = MagicMock( + get_deployment_credentials_with_provider=MagicMock(return_value=None), + afile_delete=AsyncMock(return_value=FileDeleted(id=provider_file_id, object="file", deleted=True)), + ) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert response.id == unified_file_id + assert response.object == "file" + assert response.filename == stored_file.filename + assert stored_file.id == provider_file_id + router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id) + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + @pytest.mark.asyncio async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): """ From 3b199cd3da3fc97a6a373a1247fb798fa2353ca1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:11:58 -0700 Subject: [PATCH 019/157] fix(azure_ai): price seven Foundry catalog names and charge the model router fee once Add cost map entries for azure_ai/gpt-chat-latest, codex-mini, whisper, model-router, cohere-command-a, grok-4-20-reasoning, and grok-4-20-non-reasoning, priced from the live Azure AI Foundry and Azure OpenAI pricing pages and the Azure Retail Prices API. Skip the model router flat fee when the response model is the router entry itself, since the generic cost already priced that fee. Before, azure_ai/model_router charged it twice. Resolves LIT-3157 --- litellm/llms/azure_ai/cost_calculator.py | 76 +++----- ...odel_prices_and_context_window_backup.json | 130 +++++++++++++ model_prices_and_context_window.json | 130 +++++++++++++ .../azure_ai/test_azure_ai_cost_calculator.py | 26 +++ ...azure_ai_foundry_catalog_model_metadata.py | 176 ++++++++++++++++++ 5 files changed, 492 insertions(+), 46 deletions(-) create mode 100644 tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 95f536296a2..141148f06e7 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -56,6 +56,27 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl return 0.0 +ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"}) + + +def _prices_router_fee_itself(model: str) -> bool: + return model.lower().rsplit("/", 1)[-1] in ROUTER_FEE_ENTRY_NAMES + + +def _base_cost_per_token(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float] | None: + try: + return generic_cost_per_token( + model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier + ) + except Exception as e: + if not _is_azure_model_router(model): + raise + verbose_logger.debug( + "Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e + ) + return None + + def cost_per_token( model: str, usage: Usage, @@ -66,9 +87,9 @@ def cost_per_token( """ Calculate the cost per token for Azure AI models. - For Azure AI Foundry Model Router: - - Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json) - - Plus the cost of the actual model used (handled by generic_cost_per_token) + For Azure AI Foundry Model Router the routing fee (the azure_ai/model_router entry, $0.14 per + million input tokens) is added on top of the routed model's cost. When the response model is + the router entry itself, generic_cost_per_token has already charged that fee. Args: model: str, the model name without provider prefix (from response) @@ -83,49 +104,12 @@ def cost_per_token( ValueError: If the model is not found in the cost map and cost cannot be calculated (except for Model Router models where we return just the routing flat cost) """ - prompt_cost = 0.0 - completion_cost = 0.0 - - # Determine if this was a model router request - # Check both the response model and the request model is_router_request: Final = _is_azure_model_router(model) or ( request_model is not None and _is_azure_model_router(request_model) ) - - # Calculate base cost using generic cost calculator - # This may raise an exception if the model is not in the cost map - try: - prompt_cost, completion_cost = generic_cost_per_token( - model=model, - usage=usage, - custom_llm_provider="azure_ai", - service_tier=service_tier, - ) - except Exception as e: - # For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map - # because it's a routing service, not an actual model. In this case, we continue - # to calculate just the routing flat cost. - if not _is_azure_model_router(model): - # Re-raise for non-router models - they should have pricing defined - raise - verbose_logger.debug( - "Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e - ) - - # Add flat cost for Azure Model Router - # The flat cost is defined in model_prices_and_context_window.json for azure_ai/model_router - if is_router_request: - # Use the request model for flat cost calculation if available, otherwise use response model - router_model_for_calc: Final = request_model if request_model else model - router_flat_cost: Final = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens) - - if router_flat_cost > 0: - verbose_logger.debug( - f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} " - f"({usage.prompt_tokens} tokens × ${router_flat_cost / usage.prompt_tokens:.9f}/token)" - ) - - # Add flat cost to prompt cost - prompt_cost += router_flat_cost - - return prompt_cost, completion_cost + base_cost: Final = _base_cost_per_token(model=model, usage=usage, service_tier=service_tier) + prompt_cost, completion_cost = base_cost if base_cost is not None else (0.0, 0.0) + if not is_router_request or (base_cost is not None and _prices_router_fee_itself(model)): + return prompt_cost, completion_cost + router_flat_cost: Final = calculate_azure_model_router_flat_cost(request_model or model, usage.prompt_tokens) + return prompt_cost + router_flat_cost, completion_cost diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1ffc1583e4..6649fa831d7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3581,6 +3581,82 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure_ai/codex-mini": { + "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 6e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/whisper": { + "deprecation_date": "2026-12-15", + "input_cost_per_second": 0.0001, + "litellm_provider": "azure_ai", + "mode": "audio_transcription", + "output_cost_per_second": 0.0001, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" + }, "azure_ai/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -3991,6 +4067,17 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, + "azure_ai/model-router": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", + "comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, @@ -10302,6 +10389,18 @@ "/v1/ocr" ] }, + "azure_ai/cohere-command-a": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", + "supports_function_calling": true, + "supports_tool_choice": true + }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.0015, @@ -10653,6 +10752,37 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4-20-reasoning": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_reasoning": true + }, + "azure_ai/grok-4-20-non-reasoning": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1ffc1583e4..6649fa831d7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3581,6 +3581,82 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure_ai/codex-mini": { + "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 6e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/whisper": { + "deprecation_date": "2026-12-15", + "input_cost_per_second": 0.0001, + "litellm_provider": "azure_ai", + "mode": "audio_transcription", + "output_cost_per_second": 0.0001, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" + }, "azure_ai/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -3991,6 +4067,17 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, + "azure_ai/model-router": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", + "comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, @@ -10302,6 +10389,18 @@ "/v1/ocr" ] }, + "azure_ai/cohere-command-a": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", + "supports_function_calling": true, + "supports_tool_choice": true + }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.0015, @@ -10653,6 +10752,37 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4-20-reasoning": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_reasoning": true + }, + "azure_ai/grok-4-20-non-reasoning": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 9612d97d946..80cd99bd46b 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -528,3 +528,29 @@ def test_mai_thinking_1_model_info_and_cost(local_model_cost_map): assert model_info["supports_function_calling"] is True assert prompt_cost == pytest.approx(2.0) assert completion_cost == pytest.approx(8.0) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) +def test_router_entry_as_response_model_charges_the_fee_once(router_entry_name: str) -> None: + usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) + prompt_cost, completion_cost = cost_per_token(model=router_entry_name, usage=usage) + assert prompt_cost == pytest.approx(0.14, rel=1e-9) + assert completion_cost == 0.0 + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_unmapped_router_deployment_name_still_charges_the_fee() -> None: + usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) + prompt_cost, completion_cost = cost_per_token(model="azure-model-router", usage=usage) + assert prompt_cost == pytest.approx(0.14, rel=1e-9) + assert completion_cost == 0.0 + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_routed_model_response_adds_the_fee_on_top() -> None: + usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) + routed_prompt_cost, _ = cost_per_token(model="gpt-5-nano", usage=usage) + prompt_cost, _ = cost_per_token(model="gpt-5-nano", usage=usage, request_model="azure_ai/model-router") + assert routed_prompt_cost > 0 + assert prompt_cost == pytest.approx(routed_prompt_cost + 0.14, rel=1e-9) diff --git a/tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py new file mode 100644 index 00000000000..9c5ca26a89c --- /dev/null +++ b/tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py @@ -0,0 +1,176 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm import cost_per_token, get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +REPO_ROOT: Final = Path(__file__).parents[2] +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +AZURE_OPENAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" +FOUNDRY_AOAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/" +FOUNDRY_COHERE_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/" +FOUNDRY_GROK_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/" + + +@dataclass(frozen=True, slots=True) +class TokenPricedCatalogModel: + catalog_name: str + mode: str + source: str + input_cost_per_token: float + output_cost_per_token: float + max_input_tokens: int + max_output_tokens: int + cache_read_input_token_cost: float | None + supported_flags: tuple[str, ...] + + +TOKEN_PRICED_MODELS: Final = ( + TokenPricedCatalogModel( + catalog_name="gpt-chat-latest", + mode="chat", + source=AZURE_OPENAI_PRICING, + input_cost_per_token=5e-06, + output_cost_per_token=3e-05, + max_input_tokens=200000, + max_output_tokens=128000, + cache_read_input_token_cost=5e-07, + supported_flags=( + "supports_function_calling", + "supports_prompt_caching", + "supports_reasoning", + "supports_response_schema", + "supports_tool_choice", + "supports_vision", + "supports_web_search", + ), + ), + TokenPricedCatalogModel( + catalog_name="codex-mini", + mode="responses", + source=AZURE_OPENAI_PRICING, + input_cost_per_token=1.5e-06, + output_cost_per_token=6e-06, + max_input_tokens=200000, + max_output_tokens=100000, + cache_read_input_token_cost=3.75e-07, + supported_flags=("supports_function_calling", "supports_prompt_caching", "supports_reasoning", "supports_vision"), + ), + TokenPricedCatalogModel( + catalog_name="model-router", + mode="chat", + source=FOUNDRY_AOAI_PRICING, + input_cost_per_token=1.4e-07, + output_cost_per_token=0.0, + max_input_tokens=1048576, + max_output_tokens=32768, + cache_read_input_token_cost=None, + supported_flags=(), + ), + TokenPricedCatalogModel( + catalog_name="cohere-command-a", + mode="chat", + source=FOUNDRY_COHERE_PRICING, + input_cost_per_token=2.5e-06, + output_cost_per_token=1e-05, + max_input_tokens=131072, + max_output_tokens=4096, + cache_read_input_token_cost=None, + supported_flags=("supports_function_calling", "supports_tool_choice"), + ), + TokenPricedCatalogModel( + catalog_name="grok-4-20-reasoning", + mode="chat", + source=FOUNDRY_GROK_PRICING, + input_cost_per_token=1.25e-06, + output_cost_per_token=2.5e-06, + max_input_tokens=262000, + max_output_tokens=8192, + cache_read_input_token_cost=None, + supported_flags=( + "supports_function_calling", + "supports_reasoning", + "supports_response_schema", + "supports_tool_choice", + "supports_vision", + "supports_web_search", + ), + ), + TokenPricedCatalogModel( + catalog_name="grok-4-20-non-reasoning", + mode="chat", + source=FOUNDRY_GROK_PRICING, + input_cost_per_token=1.25e-06, + output_cost_per_token=2.5e-06, + max_input_tokens=262000, + max_output_tokens=8192, + cache_read_input_token_cost=None, + supported_flags=( + "supports_function_calling", + "supports_response_schema", + "supports_tool_choice", + "supports_vision", + "supports_web_search", + ), + ), +) +CATALOG_NAMES: Final = tuple(spec.catalog_name for spec in TOKEN_PRICED_MODELS) + ("whisper",) + + +def _cost_map_entry(path: Path, catalog_name: str) -> dict[str, object]: + return COST_MAP_ADAPTER.validate_json(path.read_bytes())[f"azure_ai/{catalog_name}"] + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) +def test_azure_ai_catalog_name_is_priced_and_routed(spec: TokenPricedCatalogModel) -> None: + routed_model, provider, _, _ = get_llm_provider(model=f"azure_ai/{spec.catalog_name}") + assert (routed_model, provider) == (spec.catalog_name, "azure_ai") + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["litellm_provider"] == "azure_ai" + assert info["mode"] == spec.mode + assert info["input_cost_per_token"] == spec.input_cost_per_token + assert info["output_cost_per_token"] == spec.output_cost_per_token + assert info["cache_read_input_token_cost"] == spec.cache_read_input_token_cost + assert info["max_input_tokens"] == spec.max_input_tokens + assert info["max_output_tokens"] == spec.max_output_tokens + assert info["max_tokens"] == spec.max_output_tokens + for flag in spec.supported_flags: + assert info[flag] is True, flag + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "spec", [spec for spec in TOKEN_PRICED_MODELS if spec.catalog_name != "model-router"], ids=lambda spec: spec.catalog_name +) +def test_azure_ai_catalog_name_costs_a_million_tokens_at_list_price(spec: TokenPricedCatalogModel) -> None: + prompt_cost, completion_cost = cost_per_token( + model=f"azure_ai/{spec.catalog_name}", prompt_tokens=1_000_000, completion_tokens=1_000_000 + ) + assert prompt_cost == pytest.approx(spec.input_cost_per_token * 1_000_000) + assert completion_cost == pytest.approx(spec.output_cost_per_token * 1_000_000) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_azure_ai_whisper_catalog_name_is_priced_per_second() -> None: + routed_model, provider, _, _ = get_llm_provider(model="azure_ai/whisper") + assert (routed_model, provider) == ("whisper", "azure_ai") + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["mode"] == "audio_transcription" + assert info["input_cost_per_second"] == 0.0001 + assert info["output_cost_per_second"] == 0.0001 + + +@pytest.mark.parametrize("catalog_name", CATALOG_NAMES) +def test_azure_ai_catalog_entry_source_and_backup_match(catalog_name: str) -> None: + main_entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json", catalog_name) + backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json", catalog_name) + + assert str(main_entry["source"]).startswith("https://azure.microsoft.com/en-us/pricing/details/") + assert backup_entry == main_entry From 95402ccb711cdcd93f1c296579b18ec8bd32cab4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:36:40 -0700 Subject: [PATCH 020/157] test(azure_ai): move the Foundry catalog metadata test into the mapped azure_ai directory The new metadata test sat at the top of tests/test_litellm. The azure_ai metadata tests live in tests/test_litellm/llms/azure_ai next to the cost calculator test, so this moves it there and bumps its repo-root lookup by the two extra directory levels. No test changes. --- .../azure_ai}/test_azure_ai_foundry_catalog_model_metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tests/test_litellm/{ => llms/azure_ai}/test_azure_ai_foundry_catalog_model_metadata.py (99%) diff --git a/tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py similarity index 99% rename from tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py rename to tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py index 9c5ca26a89c..1b4e83438a6 100644 --- a/tests/test_litellm/test_azure_ai_foundry_catalog_model_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -8,7 +8,7 @@ from pydantic import TypeAdapter from litellm import cost_per_token, get_model_info from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -REPO_ROOT: Final = Path(__file__).parents[2] +REPO_ROOT: Final = Path(__file__).parents[4] COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) AZURE_OPENAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" FOUNDRY_AOAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/" From ba91588b15d308daf229f464b9aa89a0f7480afd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 7 Sep 2026 18:38:12 -0700 Subject: [PATCH 021/157] fix(proxy): price one cost estimate at one moment The totals, the per-token-type lines and the reported rates each resolved off-peak pricing on their own clock read, so a quote taken as a window opened could bill on one side of the boundary and report rates from the other. /cost/estimate now pins a billing moment for the whole quote, and every rate lookup answers for the pinned moment instead of reading the clock again Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- litellm/_internal_context.py | 24 +++++++++++ .../litellm_core_utils/llm_cost_calc/utils.py | 9 ++-- .../cost_tracking_settings.py | 43 +++++++++++-------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 42 ++++++++++++++++++ .../test_cost_tracking_settings.py | 30 +++++++++++++ 5 files changed, 126 insertions(+), 22 deletions(-) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index f856fe0f2b3..8132008731f 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current asyncio task and cannot be injected via HTTP request bodies. """ +from collections.abc import Generator +from contextlib import contextmanager from contextvars import ContextVar +from datetime import datetime, timezone from typing import Final # When True, suppresses async logging and billing for internal sub-calls # (e.g., emulated file-search steps that make nested LLM calls). is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False) + +# One request prices its totals, its per-token-type lines and the rates it reports on +# separate code paths. Each reads the clock for off-peak pricing, so without a pinned +# moment they can land on either side of a window boundary and disagree with each other. +_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None) + + +@contextmanager +def pinned_billing_time(moment: datetime) -> Generator[None]: + """Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read.""" + token: Final = _billing_time.set(moment) + try: + yield + finally: + _billing_time.reset(token) + + +def current_billing_time() -> datetime: + """The pinned billing moment, or now in UTC outside a pinned block.""" + pinned: Final = _billing_time.get() + return pinned if pinned is not None else datetime.now(timezone.utc) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 2bb15fb4c48..e2168528e6b 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -10,6 +10,7 @@ from typing import Any, Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm +from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( select_tier_for_input, @@ -306,7 +307,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(), or every window shifts by the host's offset. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc for window in windows: @@ -393,7 +394,7 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose hours apply only on its weekdays. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() reference_utc: Final = ( reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc) ) @@ -1187,7 +1188,7 @@ def generic_cost_per_token( usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 ) - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + billing_time: Final = current_time if current_time is not None else current_billing_time() ( prompt_base_cost, completion_base_cost, @@ -1379,7 +1380,7 @@ def _cost_map_billed_rates( vertex_location: str | None, current_time: datetime | None, ) -> BilledTokenRates: - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + billing_time: Final = current_time if current_time is not None else current_billing_time() ( prompt_base_cost, completion_base_cost, diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 17d82fd17e3..1faa66584d5 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -18,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel import litellm +from litellm._internal_context import current_billing_time, pinned_billing_time from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost from litellm.litellm_core_utils.llm_cost_calc.utils import get_billed_token_rates @@ -631,33 +632,39 @@ async def estimate_cost( function_id="cost-estimate", ) - # Use completion_cost which handles all the logic including margins/discounts - try: - cost_per_request: Final = completion_cost( - completion_response=mock_response, + # The totals, the per-token-type lines and the reported rates each resolve pricing on their + # own path. Pinning one moment keeps an off-peak window that opens mid-quote from splitting them. + billed_at: Final = current_billing_time() + with pinned_billing_time(billed_at): + # Use completion_cost which handles all the logic including margins/discounts + try: + cost_per_request: Final = completion_cost( + completion_response=mock_response, + model=resolved_model, + custom_llm_provider=resolved_provider, + custom_cost_per_token=resolved.custom_cost_per_token, + litellm_logging_obj=litellm_logging_obj, + ) + except Exception as e: + raise HTTPException( + status_code=404, + detail={ + "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" + }, + ) + + rates: Final = get_billed_token_rates( model=resolved_model, custom_llm_provider=resolved_provider, + usage=usage, custom_cost_per_token=resolved.custom_cost_per_token, - litellm_logging_obj=litellm_logging_obj, - ) - except Exception as e: - raise HTTPException( - status_code=404, - detail={ - "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" - }, + current_time=billed_at, ) per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown) daily: Final = per_request.times(request.num_requests_per_day) monthly: Final = per_request.times(request.num_requests_per_month) - rates: Final = get_billed_token_rates( - model=resolved_model, - custom_llm_provider=resolved_provider, - usage=usage, - custom_cost_per_token=resolved.custom_cost_per_token, - ) model_info: Final = _lookup_model_info(resolved_model) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 4de3c059e63..90178428018 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,9 +1,11 @@ import json +from datetime import datetime, timezone import pytest from fastapi.testclient import TestClient import litellm +from litellm._internal_context import pinned_billing_time from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -3981,6 +3983,46 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp assert breakdown.reasoning_cost == pytest.approx(200 * rates.output_cost_per_reasoning_token) +def test_a_pinned_billing_time_prices_the_totals_and_the_reported_rates_at_one_moment(monkeypatch): + """Totals and reported rates resolve off-peak pricing on separate paths that each read the + clock, so a window opening between the two reads used to leave them describing one request + at two different prices. Pinned, both must answer for the pinned moment.""" + monkeypatch.setitem( + litellm.model_cost, + "off-peak-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + off_peak_prompt_cost, off_peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + off_peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + with pinned_billing_time(datetime(2026, 1, 1, 12, 30, tzinfo=timezone.utc)): + peak_prompt_cost, peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + + assert off_peak_rates.input_cost_per_token == pytest.approx(1e-6) + assert peak_rates.input_cost_per_token == pytest.approx(3e-6) + assert off_peak_prompt_cost == pytest.approx(1000 * off_peak_rates.input_cost_per_token) + assert off_peak_completion_cost == pytest.approx(500 * off_peak_rates.output_cost_per_token) + assert peak_prompt_cost == pytest.approx(1000 * peak_rates.input_cost_per_token) + assert peak_completion_cost == pytest.approx(500 * peak_rates.output_cost_per_token) + + def test_billed_token_rates_are_none_for_an_unpriced_model(): usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 9847c4092df..e9485f3a044 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -4,6 +4,7 @@ Tests for cost tracking settings management endpoints. Tests the GET and PATCH endpoints for managing cost discount configuration. """ +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,6 +13,7 @@ from pydantic import ValidationError import litellm +from litellm._internal_context import pinned_billing_time from litellm.proxy._types import CostEstimateRequest from litellm.proxy.management_endpoints.cost_tracking_settings import router from litellm.proxy.proxy_server import app @@ -1108,6 +1110,34 @@ class TestEstimateCostCacheAndReasoningTokens: ) assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + @pytest.mark.asyncio + async def test_a_quote_prices_its_totals_and_its_rates_at_the_same_moment(self, monkeypatch): + """The totals and the reported rates resolve off-peak pricing on separate paths. A quote + taken as a window opens must not bill on one side of it and report rates from the other.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + response = await _estimate(None, model=A_MAPPED_MODEL) + + assert response.input_cost_per_token == pytest.approx(1e-6) + assert response.output_cost_per_token == pytest.approx(5e-6) + assert response.input_cost_per_request == pytest.approx(INPUT_TOKENS * response.input_cost_per_token) + assert response.output_cost_per_request == pytest.approx(OUTPUT_TOKENS * response.output_cost_per_token) + class TestCostEstimateRequestTokenSubsets: def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): From 415bdbfd8f6ba9bd0422087cc34509bde962ce55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 20:46:39 -0700 Subject: [PATCH 022/157] fix(azure_ai): charge the Model Router fee once and correct catalog limits The router fee was folded into azure_ai.cost_per_token and then added again by the additional_costs hook, so every routed request paid it twice. The hook now owns the fee, the entry named by the deployment supplies the price, and a response priced as the router entry itself is not charged again model-router, gpt-chat-latest and cohere-command-a carry the limits from the Foundry models page, and model-router and grok-4-20-* carry their retirement dates. The router tests now run at the completion_cost level with a Logging object, which is the path the proxy takes, and fail at the merge base --- litellm/cost_calculator.py | 9 +- litellm/llms/azure_ai/cost_calculator.py | 87 ++- ...odel_prices_and_context_window_backup.json | 11 +- model_prices_and_context_window.json | 11 +- .../azure_ai/test_azure_ai_cost_calculator.py | 561 ++++++------------ ...azure_ai_foundry_catalog_model_metadata.py | 30 +- 6 files changed, 246 insertions(+), 463 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9a9d2ceda03..fc896a098b3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -45,6 +45,9 @@ from litellm.llms.azure.cost_calculation import ( from litellm.llms.azure_ai.cost_calculator import ( cost_per_token as azure_ai_cost_per_token, ) +from litellm.llms.azure_ai.cost_calculator import ( + is_router_fee_entry as azure_ai_is_router_fee_entry, +) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.llms.bedrock.cost_calculation import ( cost_per_token as bedrock_cost_per_token, @@ -338,8 +341,6 @@ def cost_per_token( ### VERTEX LOCATION ### vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") response: Any | None = None, - ### REQUEST MODEL ### - request_model: str | None = None, # original request model for router detection ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -661,7 +662,6 @@ def cost_per_token( model=model, usage=usage_block, response_time_ms=response_time_ms, - request_model=request_model, service_tier=service_tier, ) else: @@ -1659,11 +1659,10 @@ def completion_cost( data_residency=data_residency, vertex_location=vertex_location, response=completion_response, - request_model=request_model_for_cost, ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) - if custom_llm_provider == "azure_ai": + if custom_llm_provider == "azure_ai" and not azure_ai_is_router_fee_entry(model): model_for_additional_costs = request_model_for_cost if completion_response is not None: hidden_params = getattr(completion_response, "_hidden_params", None) or {} diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 141148f06e7..8a48860c0c4 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -31,6 +31,18 @@ def _is_azure_model_router(model: str) -> bool: return "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router" +ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"}) + + +def is_router_fee_entry(model: str) -> bool: + return model.lower().removeprefix("azure_ai/") in ROUTER_FEE_ENTRY_NAMES + + +def _router_fee_entry_name(model: str) -> str: + entry_name: Final = model.lower().removeprefix("azure_ai/") + return entry_name if entry_name in ROUTER_FEE_ENTRY_NAMES else "model_router" + + def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float: """ Calculate the flat cost for Azure AI Foundry Model Router. @@ -44,26 +56,39 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl """ if not _is_azure_model_router(model): return 0.0 - - # Get the model router pricing from model_prices_and_context_window.json - # Use "model_router" as the key (without actual model name suffix) - model_info: Final = get_model_info(model="model_router", custom_llm_provider="azure_ai") + model_info: Final = get_model_info(model=_router_fee_entry_name(model), custom_llm_provider="azure_ai") router_flat_cost_per_token: Final = model_info.get("input_cost_per_token", 0) - if router_flat_cost_per_token and router_flat_cost_per_token > 0: return prompt_tokens * router_flat_cost_per_token - return 0.0 -ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"}) +def cost_per_token( + model: str, + usage: Usage, + response_time_ms: float | None = 0.0, + service_tier: str | None = None, +) -> tuple[float, float]: + """ + Price the response model's own tokens for Azure AI. + The Azure AI Foundry Model Router fee is not part of this: completion_cost charges it once through + AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost + breakdown, and a response priced as the router entry itself already carries it. A router deployment name + that is missing from the cost map prices at zero here so that line item is the whole cost. -def _prices_router_fee_itself(model: str) -> bool: - return model.lower().rsplit("/", 1)[-1] in ROUTER_FEE_ENTRY_NAMES + Args: + model: str, the model name without provider prefix (from response) + usage: LiteLLM Usage block + response_time_ms: Optional response time in milliseconds + service_tier: Optional service tier the request was priced on + Returns: + Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd -def _base_cost_per_token(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float] | None: + Raises: + ValueError: If a model that is not a Model Router name is missing from the cost map + """ try: return generic_cost_per_token( model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier @@ -72,44 +97,6 @@ def _base_cost_per_token(model: str, usage: Usage, service_tier: str | None) -> if not _is_azure_model_router(model): raise verbose_logger.debug( - "Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e + "Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e ) - return None - - -def cost_per_token( - model: str, - usage: Usage, - response_time_ms: float | None = 0.0, - request_model: str | None = None, - service_tier: str | None = None, -) -> tuple[float, float]: - """ - Calculate the cost per token for Azure AI models. - - For Azure AI Foundry Model Router the routing fee (the azure_ai/model_router entry, $0.14 per - million input tokens) is added on top of the routed model's cost. When the response model is - the router entry itself, generic_cost_per_token has already charged that fee. - - Args: - model: str, the model name without provider prefix (from response) - usage: LiteLLM Usage block - response_time_ms: Optional response time in milliseconds - request_model: Optional[str], the original request model name (to detect router usage) - - Returns: - Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd - - Raises: - ValueError: If the model is not found in the cost map and cost cannot be calculated - (except for Model Router models where we return just the routing flat cost) - """ - is_router_request: Final = _is_azure_model_router(model) or ( - request_model is not None and _is_azure_model_router(request_model) - ) - base_cost: Final = _base_cost_per_token(model=model, usage=usage, service_tier=service_tier) - prompt_cost, completion_cost = base_cost if base_cost is not None else (0.0, 0.0) - if not is_router_request or (base_cost is not None and _prices_router_fee_itself(model)): - return prompt_cost, completion_cost - router_flat_cost: Final = calculate_azure_model_router_flat_cost(request_model or model, usage.prompt_tokens) - return prompt_cost + router_flat_cost, completion_cost + return 0.0, 0.0 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6649fa831d7..483ebd2431c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3586,7 +3586,7 @@ "deprecation_date": "2026-12-02", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 200000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -4068,10 +4068,11 @@ "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure_ai/model-router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", - "max_input_tokens": 1048576, + "max_input_tokens": 200000, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", @@ -10393,8 +10394,8 @@ "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 8182, + "max_tokens": 8182, "mode": "chat", "output_cost_per_token": 1e-05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", @@ -10753,6 +10754,7 @@ "supports_web_search": true }, "azure_ai/grok-4-20-reasoning": { + "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -10769,6 +10771,7 @@ "supports_reasoning": true }, "azure_ai/grok-4-20-non-reasoning": { + "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6649fa831d7..483ebd2431c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3586,7 +3586,7 @@ "deprecation_date": "2026-12-02", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 200000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -4068,10 +4068,11 @@ "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure_ai/model-router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", - "max_input_tokens": 1048576, + "max_input_tokens": 200000, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", @@ -10393,8 +10394,8 @@ "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 8182, + "max_tokens": 8182, "mode": "chat", "output_cost_per_token": 1e-05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", @@ -10753,6 +10754,7 @@ "supports_web_search": true }, "azure_ai/grok-4-20-reasoning": { + "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -10769,6 +10771,7 @@ "supports_reasoning": true }, "azure_ai/grok-4-20-non-reasoning": { + "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 80cd99bd46b..20d0ec03a2a 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -2,20 +2,25 @@ Test Azure AI cost calculator, especially Model Router flat cost. """ +from datetime import datetime +from typing import Final + import pytest +import litellm +from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.cost_calculator import ( _is_azure_model_router, + calculate_azure_model_router_flat_cost, cost_per_token, ) -from litellm.types.utils import Usage +from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info # Get the flat cost from model_prices_and_context_window.json _model_info = get_model_info(model="model_router", custom_llm_provider="azure_ai") -AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS = ( - _model_info.get("input_cost_per_token", 0) * 1_000_000 -) +AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS = _model_info.get("input_cost_per_token", 0) * 1_000_000 class TestAzureModelRouterDetection: @@ -80,377 +85,172 @@ class TestAzureModelRouterPrefix: assert result == expected +ROUTER_FEE_PER_TOKEN: Final = AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 +ROUTED_MODEL: Final = "gpt-4.1-nano-2025-04-14" +ROUTED_USAGE: Final = Usage(prompt_tokens=5000, completion_tokens=2000, total_tokens=7000) +ROUTED_FEE: Final = 5000 * ROUTER_FEE_PER_TOKEN + + +def _router_logging(request_model: str) -> Logging: + return Logging( + model=request_model, + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-123", + function_id="test-function", + ) + + +def _azure_ai_response(response_model: str, litellm_model_name: str | None = None) -> ModelResponse: + response: Final = ModelResponse( + id="test-123", + choices=[Choices(finish_reason="stop", index=0, message=Message(role="assistant", content="Hello"))], + created=1234567890, + model=response_model, + object="chat.completion", + usage=ROUTED_USAGE, + ) + response._hidden_params = ( + {"custom_llm_provider": "azure_ai"} + if litellm_model_name is None + else {"custom_llm_provider": "azure_ai", "litellm_model_name": litellm_model_name} + ) + return response + + +def _routed_model_cost() -> tuple[float, float]: + routed_info: Final = get_model_info(model=ROUTED_MODEL, custom_llm_provider="azure_ai") + return ( + ROUTED_USAGE.prompt_tokens * (routed_info["input_cost_per_token"] or 0.0), + ROUTED_USAGE.completion_tokens * (routed_info["output_cost_per_token"] or 0.0), + ) + + +@pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterFlatCost: - """Test Azure AI Foundry Model Router flat cost calculation.""" + """cost_per_token prices the response model only; the router fee is the cost breakdown's own line item.""" - def test_model_router_flat_cost_basic(self): - """Test that flat cost is added for Model Router requests.""" - model = "azure-model-router" - usage = Usage( - prompt_tokens=1000, - completion_tokens=500, - total_tokens=1500, + def test_unmapped_router_deployment_name_prices_at_zero(self) -> None: + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + assert cost_per_token(model="azure-model-router", usage=usage) == (0.0, 0.0) + + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_router_entry_prices_its_own_fee(self, router_entry_name: str) -> None: + usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) + prompt_cost, completion_cost_usd = cost_per_token(model=router_entry_name, usage=usage) + assert prompt_cost == pytest.approx(0.14, rel=1e-9) + assert completion_cost_usd == 0.0 + + def test_routed_model_is_priced_as_itself(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = cost_per_token(model=ROUTED_MODEL, usage=ROUTED_USAGE) + assert routed_prompt_cost > 0 + assert prompt_cost == pytest.approx(routed_prompt_cost, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) + + def test_unmapped_model_that_is_not_a_router_name_raises(self) -> None: + usage = Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20) + with pytest.raises(Exception, match="no-such-azure-ai-model"): + cost_per_token(model="no-such-azure-ai-model", usage=usage) + + def test_flat_cost_helper(self) -> None: + assert calculate_azure_model_router_flat_cost( + model="azure-model-router", prompt_tokens=10_000 + ) == pytest.approx(0.0014, rel=1e-9) + assert calculate_azure_model_router_flat_cost(model="gpt-5-nano", prompt_tokens=10_000) == 0.0 + + def test_flat_cost_reads_the_fee_from_the_deployment_named_entry(self) -> None: + litellm.register_model( + {"azure_ai/model-router": {"input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "mode": "chat"}} ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 + litellm.get_model_info.cache_clear() + assert calculate_azure_model_router_flat_cost(model="model-router", prompt_tokens=1_000_000) == pytest.approx( + 0.2, rel=1e-9 ) - - # Flat cost should be $0.00014 (1000 tokens × $0.14 / 1M tokens) - assert expected_flat_cost == pytest.approx(0.00014, rel=1e-9) - - # Prompt cost should include the flat cost - # (plus any base cost from the actual model used, which might be 0 if not in model_cost) - assert prompt_cost >= expected_flat_cost - print( - f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" - ) - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_model_router_flat_cost_large_request(self): - """Test flat cost calculation for larger requests.""" - model = "model-router" - usage = Usage( - prompt_tokens=100_000, - completion_tokens=50_000, - total_tokens=150_000, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 - ) - - # Flat cost should be $0.014 (100k tokens × $0.14 / 1M tokens) - assert expected_flat_cost == pytest.approx(0.014, rel=1e-9) - # Use approx for floating-point comparison - assert prompt_cost >= expected_flat_cost or prompt_cost == pytest.approx( - expected_flat_cost, rel=1e-9 - ) - print( - f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" - ) - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_model_router_flat_cost_1m_tokens(self): - """Test flat cost for exactly 1 million input tokens.""" - model = "azure-model-router" - usage = Usage( - prompt_tokens=1_000_000, - completion_tokens=100_000, - total_tokens=1_100_000, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - - # Flat cost should be exactly $0.14 for 1M tokens - assert expected_flat_cost == pytest.approx(0.14, rel=1e-9) - assert prompt_cost >= expected_flat_cost - print(f"Model Router flat cost for 1M tokens: ${expected_flat_cost:.6f}") - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_non_model_router_no_flat_cost(self): - """Test that non-Model Router models don't get the flat cost.""" - model = "gpt-4o" - usage = Usage( - prompt_tokens=1000, - completion_tokens=500, - total_tokens=1500, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # No flat cost should be added for non-Model Router models - # The cost might be 0 or based on the model's pricing - print(f"Non-Model Router prompt cost: ${prompt_cost:.6f}") - # We just ensure it doesn't crash and returns valid values - assert prompt_cost >= 0 - assert completion_cost >= 0 - - def test_model_router_with_cached_tokens(self): - """Test Model Router flat cost with cached tokens.""" - model = "azure-model-router" - usage = Usage( - prompt_tokens=2000, - completion_tokens=800, - total_tokens=2800, - cache_read_input_tokens=500, - cache_creation_input_tokens=200, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Flat cost is based on ALL prompt tokens (including cached) - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 - ) - - assert expected_flat_cost == pytest.approx(0.00028, rel=1e-9) - assert prompt_cost >= expected_flat_cost - print( - f"Model Router flat cost with caching for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" - ) - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_router_flat_cost_when_response_has_actual_model(self): - """ - Test that router flat cost is added when request was via router but response - contains the actual model (e.g., gpt-5-nano). - - This is the key fix: Azure returns the actual model in the response, but we - must still add the router flat cost because the request was made via model router. - """ - usage = Usage( - prompt_tokens=10000, - completion_tokens=5000, - total_tokens=15000, - ) - - # Response model is the actual model Azure used (not a router name) - response_model = "gpt-5-nano-2025-08-07" - # Request model is the router - user called azure_ai/model_router/model-router - request_model = "azure_ai/model_router/model-router" - - prompt_cost, completion_cost = cost_per_token( - model=response_model, - usage=usage, - request_model=request_model, - ) - - # Expected: model cost (from gpt-5-nano) + router flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 - ) - assert expected_flat_cost == pytest.approx(0.0014, rel=1e-9) - - # Total cost should be model cost + flat cost - total_cost = prompt_cost + completion_cost - assert total_cost >= expected_flat_cost - - # Prompt cost should include both model prompt cost and router flat cost - assert prompt_cost >= expected_flat_cost + assert calculate_azure_model_router_flat_cost( + model="azure-model-router", prompt_tokens=1_000_000 + ) == pytest.approx(0.14, rel=1e-9) +@pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterCostBreakdown: - """Test that Azure Model Router flat cost is tracked in cost breakdown.""" + """completion_cost charges the router fee exactly once, as the cost breakdown's additional cost line.""" - def test_flat_cost_calculation_helper(self): - """Test that flat cost can be calculated using the helper function.""" - from litellm.llms.azure_ai.cost_calculator import ( - calculate_azure_model_router_flat_cost, - ) - - model = "azure-model-router" - prompt_tokens = 10000 - - # Calculate flat cost using helper function - flat_cost = calculate_azure_model_router_flat_cost( - model=model, prompt_tokens=prompt_tokens - ) - - # Expected flat cost - expected_flat_cost = ( - prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 - ) - - assert flat_cost > 0 - assert flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) - print(f"Flat cost calculated: ${flat_cost:.6f}") - - def test_flat_cost_integration_with_completion_cost(self): - """Test that flat cost is properly integrated into completion_cost calculation.""" - import litellm - from litellm.cost_calculator import completion_cost - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a mock response for azure_ai model router - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - role="assistant", - content="Test response", - ), - ) - ], - created=1234567890, - model="azure-model-router", - object="chat.completion", - usage=Usage( - prompt_tokens=5000, - completion_tokens=2000, - total_tokens=7000, - ), - ) - - # Set hidden params for provider - response._hidden_params = {"custom_llm_provider": "azure_ai"} - - # Calculate cost + def test_unmapped_router_deployment_name_costs_only_the_fee(self) -> None: cost = completion_cost( - completion_response=response, + completion_response=_azure_ai_response("azure-model-router"), model="azure-model-router", custom_llm_provider="azure_ai", ) + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) - # Expected flat cost - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 - ) - - # Cost should include the flat cost (use approx for floating-point comparison) - assert cost >= expected_flat_cost or cost == pytest.approx( - expected_flat_cost, rel=1e-9 - ) - print(f"Total cost with flat fee: ${cost:.6f}") - print(f"Expected minimum flat cost: ${expected_flat_cost:.6f}") - - def test_additional_costs_in_cost_breakdown(self): - """Test that Azure Model Router flat cost appears in additional_costs dict.""" - from datetime import datetime - - from litellm.cost_calculator import completion_cost - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create logging object with required parameters - logging_obj = Logging( - model="azure-model-router", - messages=[{"role": "user", "content": "Hello"}], - stream=False, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-123", - function_id="test-function", - ) - - # Create a mock response for azure_ai model router - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - role="assistant", - content="Test response", - ), - ) - ], - created=1234567890, - model="azure-model-router", - object="chat.completion", - usage=Usage( - prompt_tokens=5000, - completion_tokens=2000, - total_tokens=7000, - ), - ) - - # Set hidden params for provider - response._hidden_params = {"custom_llm_provider": "azure_ai"} - - # Calculate cost with logging object + def test_fee_is_the_breakdown_line_item_for_an_unmapped_router_name(self) -> None: + logging_obj = _router_logging("azure-model-router") cost = completion_cost( - completion_response=response, + completion_response=_azure_ai_response("azure-model-router"), model="azure-model-router", custom_llm_provider="azure_ai", litellm_logging_obj=logging_obj, ) - - # Check that cost breakdown contains additional_costs - assert hasattr(logging_obj, "cost_breakdown") - assert logging_obj.cost_breakdown is not None - assert "additional_costs" in logging_obj.cost_breakdown - assert isinstance(logging_obj.cost_breakdown["additional_costs"], dict) - - # Check that the Azure Model Router flat cost is in additional_costs - additional_costs = logging_obj.cost_breakdown["additional_costs"] - assert "Azure Model Router Flat Cost" in additional_costs - - # Verify the flat cost value - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == 0.0 + assert breakdown.get("additional_costs") == pytest.approx( + {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 ) - actual_flat_cost = additional_costs["Azure Model Router Flat Cost"] - assert actual_flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) - print(f"Additional costs in breakdown: {additional_costs}") - print(f"Azure Model Router Flat Cost: ${actual_flat_cost:.6f}") - - def test_additional_costs_when_response_has_actual_model_via_hidden_params(self): - """additional_costs populated when response has actual model but request was via model router (hidden_params).""" - from datetime import datetime - - from litellm.cost_calculator import completion_cost - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - logging_obj = Logging( - model="gpt-4.1-nano-2025-04-14", - messages=[{"role": "user", "content": "Hello"}], - stream=False, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-123", - function_id="test-function", - ) - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message(role="assistant", content="Hello"), - ) - ], - created=1234567890, - model="gpt-4.1-nano-2025-04-14", - object="chat.completion", - usage=Usage(prompt_tokens=5000, completion_tokens=2000, total_tokens=7000), - ) - response._hidden_params = { - "custom_llm_provider": "azure_ai", - "litellm_model_name": "azure_ai/model-router", - } + def test_router_request_with_routed_response_charges_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + logging_obj = _router_logging("model-router") cost = completion_cost( - completion_response=response, - model="gpt-4.1-nano-2025-04-14", + completion_response=_azure_ai_response(ROUTED_MODEL), + model=ROUTED_MODEL, custom_llm_provider="azure_ai", litellm_logging_obj=logging_obj, ) - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == pytest.approx(routed_prompt_cost, rel=1e-9) + assert breakdown["output_cost"] == pytest.approx(routed_completion_cost, rel=1e-9) + assert breakdown.get("additional_costs") == pytest.approx( + {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 ) - assert cost >= expected_flat_cost - assert logging_obj.cost_breakdown is not None - assert "additional_costs" in logging_obj.cost_breakdown - assert ( - "Azure Model Router Flat Cost" - in logging_obj.cost_breakdown["additional_costs"] + assert cost == pytest.approx(routed_prompt_cost + routed_completion_cost + ROUTED_FEE, rel=1e-9) + + def test_routed_response_named_by_hidden_params_charges_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + logging_obj = _router_logging(ROUTED_MODEL) + cost = completion_cost( + completion_response=_azure_ai_response(ROUTED_MODEL, litellm_model_name="azure_ai/model-router"), + model=ROUTED_MODEL, + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, ) - assert logging_obj.cost_breakdown["additional_costs"][ - "Azure Model Router Flat Cost" - ] == pytest.approx(expected_flat_cost, rel=1e-9) + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == pytest.approx(routed_prompt_cost, rel=1e-9) + assert breakdown.get("additional_costs") == pytest.approx( + {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 + ) + assert cost == pytest.approx(routed_prompt_cost + routed_completion_cost + ROUTED_FEE, rel=1e-9) + + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_response_priced_as_the_router_entry_charges_the_fee_once(self, router_entry_name: str) -> None: + logging_obj = _router_logging(router_entry_name) + cost = completion_cost( + completion_response=_azure_ai_response(router_entry_name), + model=router_entry_name, + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, + ) + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert "additional_costs" not in breakdown + assert breakdown["input_cost"] == pytest.approx(ROUTED_FEE, rel=1e-9) + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) class TestAzureAIServiceTierCostCalculation: @@ -459,26 +259,27 @@ class TestAzureAIServiceTierCostCalculation: @pytest.fixture(autouse=True) def register_test_model(self): import litellm - litellm.register_model(model_cost={ - "test-azure-ai-model": { - "input_cost_per_token": 0.001, - "output_cost_per_token": 0.002, - "input_cost_per_token_priority": 0.01, - "output_cost_per_token_priority": 0.02, - "input_cost_per_token_flex": 0.0005, - "output_cost_per_token_flex": 0.001, - "litellm_provider": "azure_ai", - "max_tokens": 8192, + + litellm.register_model( + model_cost={ + "test-azure-ai-model": { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure_ai", + "max_tokens": 8192, + } } - }) + ) def test_service_tier_priority_higher_cost(self): """Priority tier should cost more than standard for azure_ai.""" usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - standard_prompt, standard_completion = cost_per_token( - model="test-azure-ai-model", usage=usage - ) + standard_prompt, standard_completion = cost_per_token(model="test-azure-ai-model", usage=usage) priority_prompt, priority_completion = cost_per_token( model="test-azure-ai-model", usage=usage, service_tier="priority" ) @@ -490,12 +291,8 @@ class TestAzureAIServiceTierCostCalculation: """Flex tier should cost less than standard for azure_ai.""" usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - standard_prompt, standard_completion = cost_per_token( - model="test-azure-ai-model", usage=usage - ) - flex_prompt, flex_completion = cost_per_token( - model="test-azure-ai-model", usage=usage, service_tier="flex" - ) + standard_prompt, standard_completion = cost_per_token(model="test-azure-ai-model", usage=usage) + flex_prompt, flex_completion = cost_per_token(model="test-azure-ai-model", usage=usage, service_tier="flex") assert flex_prompt < standard_prompt assert flex_completion < standard_completion @@ -528,29 +325,3 @@ def test_mai_thinking_1_model_info_and_cost(local_model_cost_map): assert model_info["supports_function_calling"] is True assert prompt_cost == pytest.approx(2.0) assert completion_cost == pytest.approx(8.0) - - -@pytest.mark.usefixtures("local_model_cost_map") -@pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) -def test_router_entry_as_response_model_charges_the_fee_once(router_entry_name: str) -> None: - usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) - prompt_cost, completion_cost = cost_per_token(model=router_entry_name, usage=usage) - assert prompt_cost == pytest.approx(0.14, rel=1e-9) - assert completion_cost == 0.0 - - -@pytest.mark.usefixtures("local_model_cost_map") -def test_unmapped_router_deployment_name_still_charges_the_fee() -> None: - usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) - prompt_cost, completion_cost = cost_per_token(model="azure-model-router", usage=usage) - assert prompt_cost == pytest.approx(0.14, rel=1e-9) - assert completion_cost == 0.0 - - -@pytest.mark.usefixtures("local_model_cost_map") -def test_routed_model_response_adds_the_fee_on_top() -> None: - usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) - routed_prompt_cost, _ = cost_per_token(model="gpt-5-nano", usage=usage) - prompt_cost, _ = cost_per_token(model="gpt-5-nano", usage=usage, request_model="azure_ai/model-router") - assert routed_prompt_cost > 0 - assert prompt_cost == pytest.approx(routed_prompt_cost + 0.14, rel=1e-9) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py index 1b4e83438a6..fab9be1b42c 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -26,6 +26,7 @@ class TokenPricedCatalogModel: max_input_tokens: int max_output_tokens: int cache_read_input_token_cost: float | None + deprecation_date: str | None supported_flags: tuple[str, ...] @@ -36,9 +37,10 @@ TOKEN_PRICED_MODELS: Final = ( source=AZURE_OPENAI_PRICING, input_cost_per_token=5e-06, output_cost_per_token=3e-05, - max_input_tokens=200000, + max_input_tokens=272000, max_output_tokens=128000, cache_read_input_token_cost=5e-07, + deprecation_date="2026-12-02", supported_flags=( "supports_function_calling", "supports_prompt_caching", @@ -58,7 +60,13 @@ TOKEN_PRICED_MODELS: Final = ( max_input_tokens=200000, max_output_tokens=100000, cache_read_input_token_cost=3.75e-07, - supported_flags=("supports_function_calling", "supports_prompt_caching", "supports_reasoning", "supports_vision"), + deprecation_date="2026-11-15", + supported_flags=( + "supports_function_calling", + "supports_prompt_caching", + "supports_reasoning", + "supports_vision", + ), ), TokenPricedCatalogModel( catalog_name="model-router", @@ -66,9 +74,10 @@ TOKEN_PRICED_MODELS: Final = ( source=FOUNDRY_AOAI_PRICING, input_cost_per_token=1.4e-07, output_cost_per_token=0.0, - max_input_tokens=1048576, + max_input_tokens=200000, max_output_tokens=32768, cache_read_input_token_cost=None, + deprecation_date="2027-05-20", supported_flags=(), ), TokenPricedCatalogModel( @@ -78,8 +87,9 @@ TOKEN_PRICED_MODELS: Final = ( input_cost_per_token=2.5e-06, output_cost_per_token=1e-05, max_input_tokens=131072, - max_output_tokens=4096, + max_output_tokens=8182, cache_read_input_token_cost=None, + deprecation_date=None, supported_flags=("supports_function_calling", "supports_tool_choice"), ), TokenPricedCatalogModel( @@ -91,6 +101,7 @@ TOKEN_PRICED_MODELS: Final = ( max_input_tokens=262000, max_output_tokens=8192, cache_read_input_token_cost=None, + deprecation_date="2027-04-06", supported_flags=( "supports_function_calling", "supports_reasoning", @@ -109,6 +120,7 @@ TOKEN_PRICED_MODELS: Final = ( max_input_tokens=262000, max_output_tokens=8192, cache_read_input_token_cost=None, + deprecation_date="2027-04-06", supported_flags=( "supports_function_calling", "supports_response_schema", @@ -146,7 +158,9 @@ def test_azure_ai_catalog_name_is_priced_and_routed(spec: TokenPricedCatalogMode @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( - "spec", [spec for spec in TOKEN_PRICED_MODELS if spec.catalog_name != "model-router"], ids=lambda spec: spec.catalog_name + "spec", + [spec for spec in TOKEN_PRICED_MODELS if spec.catalog_name != "model-router"], + ids=lambda spec: spec.catalog_name, ) def test_azure_ai_catalog_name_costs_a_million_tokens_at_list_price(spec: TokenPricedCatalogModel) -> None: prompt_cost, completion_cost = cost_per_token( @@ -174,3 +188,9 @@ def test_azure_ai_catalog_entry_source_and_backup_match(catalog_name: str) -> No assert str(main_entry["source"]).startswith("https://azure.microsoft.com/en-us/pricing/details/") assert backup_entry == main_entry + + +@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) +def test_azure_ai_catalog_entry_carries_its_retirement_date(spec: TokenPricedCatalogModel) -> None: + entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json", spec.catalog_name) + assert entry.get("deprecation_date") == spec.deprecation_date From c02f2dc0feff1f95d61b1be699565a835402a122 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:10:54 -0700 Subject: [PATCH 023/157] fix(azure_ai): keep the request_model keyword on cost_per_token Restores the public keyword removed at 415bdbfd8f. A direct caller that names the Model Router as the request model gets the routing fee folded into the prompt cost once; completion_cost never passes it and charges the fee through the additional-costs hook as before --- litellm/cost_calculator.py | 3 + litellm/llms/azure_ai/cost_calculator.py | 62 +++++++++++-------- .../azure_ai/test_azure_ai_cost_calculator.py | 33 ++++++++++ 3 files changed, 72 insertions(+), 26 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index fc896a098b3..e135503d11d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -341,6 +341,8 @@ def cost_per_token( ### VERTEX LOCATION ### vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") response: Any | None = None, + ### REQUEST MODEL ### + request_model: str | None = None, # original request model for router detection ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -662,6 +664,7 @@ def cost_per_token( model=model, usage=usage_block, response_time_ms=response_time_ms, + request_model=request_model, service_tier=service_tier, ) else: diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 8a48860c0c4..e57ba055587 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -63,32 +63,7 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl return 0.0 -def cost_per_token( - model: str, - usage: Usage, - response_time_ms: float | None = 0.0, - service_tier: str | None = None, -) -> tuple[float, float]: - """ - Price the response model's own tokens for Azure AI. - - The Azure AI Foundry Model Router fee is not part of this: completion_cost charges it once through - AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost - breakdown, and a response priced as the router entry itself already carries it. A router deployment name - that is missing from the cost map prices at zero here so that line item is the whole cost. - - Args: - model: str, the model name without provider prefix (from response) - usage: LiteLLM Usage block - response_time_ms: Optional response time in milliseconds - service_tier: Optional service tier the request was priced on - - Returns: - Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd - - Raises: - ValueError: If a model that is not a Model Router name is missing from the cost map - """ +def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float]: try: return generic_cost_per_token( model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier @@ -100,3 +75,38 @@ def cost_per_token( "Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e ) return 0.0, 0.0 + + +def cost_per_token( + model: str, + usage: Usage, + response_time_ms: float | None = 0.0, + request_model: str | None = None, + service_tier: str | None = None, +) -> tuple[float, float]: + """ + Price the response model's own tokens for Azure AI, plus the Model Router fee when the caller names the + router as the request model. + + completion_cost never passes request_model: it charges the fee once through + AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost + breakdown. A response priced as the router entry itself already carries the fee, so request_model adds + nothing on top of it, and a router deployment name that is missing from the cost map prices at zero here. + + Args: + model: str, the model name without provider prefix (from response) + usage: LiteLLM Usage block + response_time_ms: Optional response time in milliseconds + request_model: Optional[str], the original request model name; a Model Router name adds the routing fee + service_tier: Optional service tier the request was priced on + + Returns: + Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + + Raises: + ValueError: If a model that is not a Model Router name is missing from the cost map + """ + prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier) + if request_model is None or not _is_azure_model_router(request_model) or is_router_fee_entry(model): + return prompt_cost, completion_cost + return prompt_cost + calculate_azure_model_router_flat_cost(request_model, usage.prompt_tokens), completion_cost diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 20d0ec03a2a..0deb79d14d3 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -155,6 +155,39 @@ class TestAzureModelRouterFlatCost: with pytest.raises(Exception, match="no-such-azure-ai-model"): cost_per_token(model="no-such-azure-ai-model", usage=usage) + def test_request_model_through_the_router_adds_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = cost_per_token( + model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model="azure_ai/model-router" + ) + assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) + + def test_request_model_that_is_not_the_router_adds_nothing(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + assert cost_per_token( + model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model=f"azure_ai/{ROUTED_MODEL}" + ) == pytest.approx((routed_prompt_cost, routed_completion_cost), rel=1e-9) + + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_request_model_does_not_double_the_router_entry(self, router_entry_name: str) -> None: + prompt_cost, completion_cost_usd = cost_per_token( + model=router_entry_name, usage=ROUTED_USAGE, request_model=f"azure_ai/{router_entry_name}" + ) + assert prompt_cost == pytest.approx(ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == 0.0 + + def test_public_cost_per_token_keeps_the_request_model_keyword(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = litellm.cost_per_token( + model=ROUTED_MODEL, + custom_llm_provider="azure_ai", + usage_object=ROUTED_USAGE, + request_model="azure_ai/model-router", + ) + assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) + def test_flat_cost_helper(self) -> None: assert calculate_azure_model_router_flat_cost( model="azure-model-router", prompt_tokens=10_000 From 55c10c1983c92dbad1d62dfc3c14aab94696639a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:34:34 -0700 Subject: [PATCH 024/157] fix(azure_ai): charge the router fee once for any router name and price grok-4-20 cache reads Direct litellm.cost_per_token callers that name a Model Router deployment as the model get the routing fee again, as they did before this branch, and the fee is still charged exactly once on every completion_cost path. The grok-4-20 entries bill cached prompt tokens at the input rate, since Azure has no cached-input meter for them, and the model_router twin carries the same limits and retirement date as model-router. The catalog test now exercises the cost calculator and map relations instead of pinning map fields. --- litellm/cost_calculator.py | 4 +- litellm/llms/azure_ai/cost_calculator.py | 33 ++- ...odel_prices_and_context_window_backup.json | 6 + model_prices_and_context_window.json | 6 + .../azure_ai/test_azure_ai_cost_calculator.py | 31 ++- ...azure_ai_foundry_catalog_model_metadata.py | 219 ++++++------------ 6 files changed, 125 insertions(+), 174 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e135503d11d..8a00ffa4d37 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -46,7 +46,7 @@ from litellm.llms.azure_ai.cost_calculator import ( cost_per_token as azure_ai_cost_per_token, ) from litellm.llms.azure_ai.cost_calculator import ( - is_router_fee_entry as azure_ai_is_router_fee_entry, + is_azure_model_router as azure_ai_is_model_router_name, ) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.llms.bedrock.cost_calculation import ( @@ -1665,7 +1665,7 @@ def completion_cost( ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) - if custom_llm_provider == "azure_ai" and not azure_ai_is_router_fee_entry(model): + if custom_llm_provider == "azure_ai" and not azure_ai_is_model_router_name(model): model_for_additional_costs = request_model_for_cost if completion_response is not None: hidden_params = getattr(completion_response, "_hidden_params", None) or {} diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index e57ba055587..5934525eca3 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -11,7 +11,7 @@ from litellm.types.utils import Usage from litellm.utils import get_model_info -def _is_azure_model_router(model: str) -> bool: +def is_azure_model_router(model: str) -> bool: """ Check if the model is Azure AI Foundry Model Router. @@ -54,7 +54,7 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl Returns: float: The flat cost in USD, or 0.0 if not applicable """ - if not _is_azure_model_router(model): + if not is_azure_model_router(model): return 0.0 model_info: Final = get_model_info(model=_router_fee_entry_name(model), custom_llm_provider="azure_ai") router_flat_cost_per_token: Final = model_info.get("input_cost_per_token", 0) @@ -69,7 +69,7 @@ def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier ) except Exception as e: - if not _is_azure_model_router(model): + if not is_azure_model_router(model): raise verbose_logger.debug( "Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e @@ -77,6 +77,16 @@ def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> return 0.0, 0.0 +def _router_fee_name(model: str, request_model: str | None) -> str | None: + if is_router_fee_entry(model): + return None + if is_azure_model_router(model): + return model + if request_model is not None and is_azure_model_router(request_model): + return request_model + return None + + def cost_per_token( model: str, usage: Usage, @@ -85,13 +95,15 @@ def cost_per_token( service_tier: str | None = None, ) -> tuple[float, float]: """ - Price the response model's own tokens for Azure AI, plus the Model Router fee when the caller names the - router as the request model. + Price the response model's own tokens for Azure AI, plus the Model Router fee exactly once when either the + priced name or request_model is a Model Router name. - completion_cost never passes request_model: it charges the fee once through + A response priced as the router entry itself already carries the fee, so nothing is added on top of it. A + router deployment name that is missing from the cost map prices at the fee alone. + + completion_cost passes only the priced name: when that name is a routed model it adds the fee itself through AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost - breakdown. A response priced as the router entry itself already carries the fee, so request_model adds - nothing on top of it, and a router deployment name that is missing from the cost map prices at zero here. + breakdown, and when the name is router-shaped the fee is already in the prompt cost returned here. Args: model: str, the model name without provider prefix (from response) @@ -107,6 +119,7 @@ def cost_per_token( ValueError: If a model that is not a Model Router name is missing from the cost map """ prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier) - if request_model is None or not _is_azure_model_router(request_model) or is_router_fee_entry(model): + fee_name: Final = _router_fee_name(model=model, request_model=request_model) + if fee_name is None: return prompt_cost, completion_cost - return prompt_cost + calculate_azure_model_router_flat_cost(request_model, usage.prompt_tokens), completion_cost + return prompt_cost + calculate_azure_model_router_flat_cost(fee_name, usage.prompt_tokens), completion_cost diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 483ebd2431c..674a8b98304 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4060,9 +4060,13 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/model_router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" @@ -10754,6 +10758,7 @@ "supports_web_search": true }, "azure_ai/grok-4-20-reasoning": { + "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", @@ -10771,6 +10776,7 @@ "supports_reasoning": true }, "azure_ai/grok-4-20-non-reasoning": { + "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 483ebd2431c..674a8b98304 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4060,9 +4060,13 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/model_router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" @@ -10754,6 +10758,7 @@ "supports_web_search": true }, "azure_ai/grok-4-20-reasoning": { + "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", @@ -10771,6 +10776,7 @@ "supports_reasoning": true }, "azure_ai/grok-4-20-non-reasoning": { + "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-04-06", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure_ai", diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 0deb79d14d3..7df14b91741 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -11,9 +11,9 @@ import litellm from litellm.cost_calculator import completion_cost from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.cost_calculator import ( - _is_azure_model_router, calculate_azure_model_router_flat_cost, cost_per_token, + is_azure_model_router, ) from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info @@ -54,7 +54,7 @@ class TestAzureModelRouterDetection: ) def test_is_azure_model_router(self, model: str, expected: bool): """Test Azure Model Router detection.""" - assert _is_azure_model_router(model) == expected + assert is_azure_model_router(model) == expected class TestAzureModelRouterPrefix: @@ -130,11 +130,21 @@ def _routed_model_cost() -> tuple[float, float]: @pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterFlatCost: - """cost_per_token prices the response model only; the router fee is the cost breakdown's own line item.""" + """cost_per_token charges the router fee once, for whichever router name the caller gives it.""" - def test_unmapped_router_deployment_name_prices_at_zero(self) -> None: + def test_unmapped_router_deployment_name_prices_the_fee(self) -> None: usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - assert cost_per_token(model="azure-model-router", usage=usage) == (0.0, 0.0) + prompt_cost, completion_cost_usd = cost_per_token(model="azure-model-router", usage=usage) + assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 + + def test_router_deployment_name_as_both_names_charges_the_fee_once(self) -> None: + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + prompt_cost, completion_cost_usd = cost_per_token( + model="model_router/my-deployment", usage=usage, request_model="azure_ai/model_router/my-deployment" + ) + assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) def test_router_entry_prices_its_own_fee(self, router_entry_name: str) -> None: @@ -209,7 +219,8 @@ class TestAzureModelRouterFlatCost: @pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterCostBreakdown: - """completion_cost charges the router fee exactly once, as the cost breakdown's additional cost line.""" + """completion_cost charges the router fee exactly once: as the breakdown's additional cost line when a routed + model is priced as itself, inside the input cost when the priced name is the router.""" def test_unmapped_router_deployment_name_costs_only_the_fee(self) -> None: cost = completion_cost( @@ -219,7 +230,7 @@ class TestAzureModelRouterCostBreakdown: ) assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) - def test_fee_is_the_breakdown_line_item_for_an_unmapped_router_name(self) -> None: + def test_unmapped_router_name_carries_the_fee_as_its_input_cost(self) -> None: logging_obj = _router_logging("azure-model-router") cost = completion_cost( completion_response=_azure_ai_response("azure-model-router"), @@ -229,10 +240,8 @@ class TestAzureModelRouterCostBreakdown: ) breakdown = logging_obj.cost_breakdown assert breakdown is not None - assert breakdown["input_cost"] == 0.0 - assert breakdown.get("additional_costs") == pytest.approx( - {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 - ) + assert breakdown["input_cost"] == pytest.approx(ROUTED_FEE, rel=1e-9) + assert "additional_costs" not in breakdown assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) def test_router_request_with_routed_response_charges_the_fee_once(self) -> None: diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py index fab9be1b42c..19b082edd8a 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -5,131 +5,34 @@ from typing import Final import pytest from pydantic import TypeAdapter -from litellm import cost_per_token, get_model_info +from litellm import completion_cost, cost_per_token from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import TranscriptionResponse REPO_ROOT: Final = Path(__file__).parents[4] +MAIN_COST_MAP: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) -AZURE_OPENAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" -FOUNDRY_AOAI_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/" -FOUNDRY_COHERE_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/" -FOUNDRY_GROK_PRICING: Final = "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/" +AZURE_PRICING_PREFIX: Final = "https://azure.microsoft.com/en-us/pricing/details/" +A_MILLION: Final = 1_000_000 @dataclass(frozen=True, slots=True) class TokenPricedCatalogModel: catalog_name: str - mode: str - source: str - input_cost_per_token: float - output_cost_per_token: float - max_input_tokens: int - max_output_tokens: int - cache_read_input_token_cost: float | None - deprecation_date: str | None - supported_flags: tuple[str, ...] + dollars_per_million_input: float + dollars_per_million_output: float TOKEN_PRICED_MODELS: Final = ( - TokenPricedCatalogModel( - catalog_name="gpt-chat-latest", - mode="chat", - source=AZURE_OPENAI_PRICING, - input_cost_per_token=5e-06, - output_cost_per_token=3e-05, - max_input_tokens=272000, - max_output_tokens=128000, - cache_read_input_token_cost=5e-07, - deprecation_date="2026-12-02", - supported_flags=( - "supports_function_calling", - "supports_prompt_caching", - "supports_reasoning", - "supports_response_schema", - "supports_tool_choice", - "supports_vision", - "supports_web_search", - ), - ), - TokenPricedCatalogModel( - catalog_name="codex-mini", - mode="responses", - source=AZURE_OPENAI_PRICING, - input_cost_per_token=1.5e-06, - output_cost_per_token=6e-06, - max_input_tokens=200000, - max_output_tokens=100000, - cache_read_input_token_cost=3.75e-07, - deprecation_date="2026-11-15", - supported_flags=( - "supports_function_calling", - "supports_prompt_caching", - "supports_reasoning", - "supports_vision", - ), - ), - TokenPricedCatalogModel( - catalog_name="model-router", - mode="chat", - source=FOUNDRY_AOAI_PRICING, - input_cost_per_token=1.4e-07, - output_cost_per_token=0.0, - max_input_tokens=200000, - max_output_tokens=32768, - cache_read_input_token_cost=None, - deprecation_date="2027-05-20", - supported_flags=(), - ), - TokenPricedCatalogModel( - catalog_name="cohere-command-a", - mode="chat", - source=FOUNDRY_COHERE_PRICING, - input_cost_per_token=2.5e-06, - output_cost_per_token=1e-05, - max_input_tokens=131072, - max_output_tokens=8182, - cache_read_input_token_cost=None, - deprecation_date=None, - supported_flags=("supports_function_calling", "supports_tool_choice"), - ), - TokenPricedCatalogModel( - catalog_name="grok-4-20-reasoning", - mode="chat", - source=FOUNDRY_GROK_PRICING, - input_cost_per_token=1.25e-06, - output_cost_per_token=2.5e-06, - max_input_tokens=262000, - max_output_tokens=8192, - cache_read_input_token_cost=None, - deprecation_date="2027-04-06", - supported_flags=( - "supports_function_calling", - "supports_reasoning", - "supports_response_schema", - "supports_tool_choice", - "supports_vision", - "supports_web_search", - ), - ), - TokenPricedCatalogModel( - catalog_name="grok-4-20-non-reasoning", - mode="chat", - source=FOUNDRY_GROK_PRICING, - input_cost_per_token=1.25e-06, - output_cost_per_token=2.5e-06, - max_input_tokens=262000, - max_output_tokens=8192, - cache_read_input_token_cost=None, - deprecation_date="2027-04-06", - supported_flags=( - "supports_function_calling", - "supports_response_schema", - "supports_tool_choice", - "supports_vision", - "supports_web_search", - ), - ), + TokenPricedCatalogModel("gpt-chat-latest", 5.0, 30.0), + TokenPricedCatalogModel("codex-mini", 1.5, 6.0), + TokenPricedCatalogModel("model-router", 0.14, 0.0), + TokenPricedCatalogModel("cohere-command-a", 2.5, 10.0), + TokenPricedCatalogModel("grok-4-20-reasoning", 1.25, 2.5), + TokenPricedCatalogModel("grok-4-20-non-reasoning", 1.25, 2.5), ) +GROK_4_20_NAMES: Final = ("grok-4-20-reasoning", "grok-4-20-non-reasoning") CATALOG_NAMES: Final = tuple(spec.catalog_name for spec in TOKEN_PRICED_MODELS) + ("whisper",) @@ -137,60 +40,74 @@ def _cost_map_entry(path: Path, catalog_name: str) -> dict[str, object]: return COST_MAP_ADAPTER.validate_json(path.read_bytes())[f"azure_ai/{catalog_name}"] +@pytest.mark.parametrize("catalog_name", CATALOG_NAMES) +def test_azure_ai_catalog_name_routes_to_azure_ai(catalog_name: str) -> None: + routed_model, provider, _, _ = get_llm_provider(model=f"azure_ai/{catalog_name}") + assert (routed_model, provider) == (catalog_name, "azure_ai") + + @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) -def test_azure_ai_catalog_name_is_priced_and_routed(spec: TokenPricedCatalogModel) -> None: - routed_model, provider, _, _ = get_llm_provider(model=f"azure_ai/{spec.catalog_name}") - assert (routed_model, provider) == (spec.catalog_name, "azure_ai") - - info = get_model_info(model=routed_model, custom_llm_provider=provider) - assert info["litellm_provider"] == "azure_ai" - assert info["mode"] == spec.mode - assert info["input_cost_per_token"] == spec.input_cost_per_token - assert info["output_cost_per_token"] == spec.output_cost_per_token - assert info["cache_read_input_token_cost"] == spec.cache_read_input_token_cost - assert info["max_input_tokens"] == spec.max_input_tokens - assert info["max_output_tokens"] == spec.max_output_tokens - assert info["max_tokens"] == spec.max_output_tokens - for flag in spec.supported_flags: - assert info[flag] is True, flag +def test_azure_ai_catalog_name_costs_a_million_tokens_at_list_price(spec: TokenPricedCatalogModel) -> None: + prompt_cost, completion_cost_usd = cost_per_token( + model=f"azure_ai/{spec.catalog_name}", prompt_tokens=A_MILLION, completion_tokens=A_MILLION + ) + assert prompt_cost == pytest.approx(spec.dollars_per_million_input) + assert completion_cost_usd == pytest.approx(spec.dollars_per_million_output) @pytest.mark.usefixtures("local_model_cost_map") -@pytest.mark.parametrize( - "spec", - [spec for spec in TOKEN_PRICED_MODELS if spec.catalog_name != "model-router"], - ids=lambda spec: spec.catalog_name, -) -def test_azure_ai_catalog_name_costs_a_million_tokens_at_list_price(spec: TokenPricedCatalogModel) -> None: - prompt_cost, completion_cost = cost_per_token( - model=f"azure_ai/{spec.catalog_name}", prompt_tokens=1_000_000, completion_tokens=1_000_000 +@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) +def test_azure_ai_catalog_name_prices_the_same_in_any_casing(spec: TokenPricedCatalogModel) -> None: + lowercase_cost = cost_per_token(model=f"azure_ai/{spec.catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + upper_cost = cost_per_token(model=f"azure_ai/{spec.catalog_name.upper()}", prompt_tokens=A_MILLION, completion_tokens=0) + assert upper_cost == lowercase_cost + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("catalog_name", GROK_4_20_NAMES) +def test_azure_ai_grok_4_20_bills_cached_prompt_tokens_at_the_input_price(catalog_name: str) -> None: + uncached_prompt_cost, _ = cost_per_token(model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + cached_prompt_cost, _ = cost_per_token( + model=f"azure_ai/{catalog_name}", + prompt_tokens=A_MILLION, + completion_tokens=0, + cache_read_input_tokens=A_MILLION, ) - assert prompt_cost == pytest.approx(spec.input_cost_per_token * 1_000_000) - assert completion_cost == pytest.approx(spec.output_cost_per_token * 1_000_000) + assert uncached_prompt_cost > 0 + assert cached_prompt_cost == pytest.approx(uncached_prompt_cost) @pytest.mark.usefixtures("local_model_cost_map") def test_azure_ai_whisper_catalog_name_is_priced_per_second() -> None: - routed_model, provider, _, _ = get_llm_provider(model="azure_ai/whisper") - assert (routed_model, provider) == ("whisper", "azure_ai") - - info = get_model_info(model=routed_model, custom_llm_provider=provider) - assert info["mode"] == "audio_transcription" - assert info["input_cost_per_second"] == 0.0001 - assert info["output_cost_per_second"] == 0.0001 + transcription: Final = TranscriptionResponse(text="hello") + transcription._hidden_params = { # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter + "custom_llm_provider": "azure_ai", + "model": "azure_ai/whisper", + "audio_transcription_duration": 3600, + } + cost = completion_cost( + completion_response=transcription, + model="azure_ai/whisper", + custom_llm_provider="azure_ai", + call_type="atranscription", + ) + assert cost == pytest.approx(0.36) @pytest.mark.parametrize("catalog_name", CATALOG_NAMES) def test_azure_ai_catalog_entry_source_and_backup_match(catalog_name: str) -> None: - main_entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json", catalog_name) - backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json", catalog_name) + main_entry = _cost_map_entry(MAIN_COST_MAP, catalog_name) + backup_entry = _cost_map_entry(BACKUP_COST_MAP, catalog_name) - assert str(main_entry["source"]).startswith("https://azure.microsoft.com/en-us/pricing/details/") + assert str(main_entry["source"]).startswith(AZURE_PRICING_PREFIX) assert backup_entry == main_entry -@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) -def test_azure_ai_catalog_entry_carries_its_retirement_date(spec: TokenPricedCatalogModel) -> None: - entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json", spec.catalog_name) - assert entry.get("deprecation_date") == spec.deprecation_date +def test_azure_ai_model_router_spellings_share_one_entry() -> None: + underscore_entry = _cost_map_entry(MAIN_COST_MAP, "model_router") + hyphen_entry = _cost_map_entry(MAIN_COST_MAP, "model-router") + + assert {k: v for k, v in underscore_entry.items() if k != "comment"} == { + k: v for k, v in hyphen_entry.items() if k != "comment" + } From 5706952588ee2b2445e864ce8a85af3339bb138b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:22:51 -0700 Subject: [PATCH 025/157] fix(azure_ai): drop gpt-chat-latest effort levels, test prices via calculator litellm's azure_ai config rejects reasoning_effort for gpt-chat-latest and Azure documents a fixed reasoning level for it, so the entry no longer advertises reasoning_effort_levels. The catalog metadata tests compare cost_per_token and the whisper transcription cost with the entry the calculator read instead of with list-price literals, the pattern #40195 removed --- ...odel_prices_and_context_window_backup.json | 3 - model_prices_and_context_window.json | 3 - ...azure_ai_foundry_catalog_model_metadata.py | 80 +++++++++---------- 3 files changed, 40 insertions(+), 46 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 674a8b98304..7e7d8a9e930 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3591,9 +3591,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, - "reasoning_effort_levels": [ - "medium" - ], "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 674a8b98304..7e7d8a9e930 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3591,9 +3591,6 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, - "reasoning_effort_levels": [ - "medium" - ], "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py index 19b082edd8a..84d5cd2a7d4 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -1,11 +1,10 @@ -from dataclasses import dataclass from pathlib import Path from typing import Final import pytest from pydantic import TypeAdapter -from litellm import completion_cost, cost_per_token +from litellm import completion_cost, cost_per_token, get_model_info from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.types.utils import TranscriptionResponse @@ -15,31 +14,39 @@ BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_windo COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) AZURE_PRICING_PREFIX: Final = "https://azure.microsoft.com/en-us/pricing/details/" A_MILLION: Final = 1_000_000 +AN_HOUR_IN_SECONDS: Final = 3600 - -@dataclass(frozen=True, slots=True) -class TokenPricedCatalogModel: - catalog_name: str - dollars_per_million_input: float - dollars_per_million_output: float - - -TOKEN_PRICED_MODELS: Final = ( - TokenPricedCatalogModel("gpt-chat-latest", 5.0, 30.0), - TokenPricedCatalogModel("codex-mini", 1.5, 6.0), - TokenPricedCatalogModel("model-router", 0.14, 0.0), - TokenPricedCatalogModel("cohere-command-a", 2.5, 10.0), - TokenPricedCatalogModel("grok-4-20-reasoning", 1.25, 2.5), - TokenPricedCatalogModel("grok-4-20-non-reasoning", 1.25, 2.5), +TOKEN_PRICED_NAMES: Final = ( + "gpt-chat-latest", + "codex-mini", + "model-router", + "cohere-command-a", + "grok-4-20-reasoning", + "grok-4-20-non-reasoning", ) GROK_4_20_NAMES: Final = ("grok-4-20-reasoning", "grok-4-20-non-reasoning") -CATALOG_NAMES: Final = tuple(spec.catalog_name for spec in TOKEN_PRICED_MODELS) + ("whisper",) +CATALOG_NAMES: Final = TOKEN_PRICED_NAMES + ("whisper",) def _cost_map_entry(path: Path, catalog_name: str) -> dict[str, object]: return COST_MAP_ADAPTER.validate_json(path.read_bytes())[f"azure_ai/{catalog_name}"] +def _whisper_transcription_cost(duration_seconds: int) -> float: + transcription: Final = TranscriptionResponse(text="hello") + transcription._hidden_params = { # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter + "custom_llm_provider": "azure_ai", + "model": "azure_ai/whisper", + "audio_transcription_duration": duration_seconds, + } + return completion_cost( + completion_response=transcription, + model="azure_ai/whisper", + custom_llm_provider="azure_ai", + call_type="atranscription", + ) + + @pytest.mark.parametrize("catalog_name", CATALOG_NAMES) def test_azure_ai_catalog_name_routes_to_azure_ai(catalog_name: str) -> None: routed_model, provider, _, _ = get_llm_provider(model=f"azure_ai/{catalog_name}") @@ -47,20 +54,22 @@ def test_azure_ai_catalog_name_routes_to_azure_ai(catalog_name: str) -> None: @pytest.mark.usefixtures("local_model_cost_map") -@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) -def test_azure_ai_catalog_name_costs_a_million_tokens_at_list_price(spec: TokenPricedCatalogModel) -> None: +@pytest.mark.parametrize("catalog_name", TOKEN_PRICED_NAMES) +def test_azure_ai_catalog_name_charges_its_own_entry_per_token(catalog_name: str) -> None: + entry: Final = get_model_info(f"azure_ai/{catalog_name}") prompt_cost, completion_cost_usd = cost_per_token( - model=f"azure_ai/{spec.catalog_name}", prompt_tokens=A_MILLION, completion_tokens=A_MILLION + model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=A_MILLION ) - assert prompt_cost == pytest.approx(spec.dollars_per_million_input) - assert completion_cost_usd == pytest.approx(spec.dollars_per_million_output) + assert prompt_cost > 0 + assert prompt_cost == pytest.approx(A_MILLION * entry["input_cost_per_token"]) + assert completion_cost_usd == pytest.approx(A_MILLION * entry["output_cost_per_token"]) @pytest.mark.usefixtures("local_model_cost_map") -@pytest.mark.parametrize("spec", TOKEN_PRICED_MODELS, ids=lambda spec: spec.catalog_name) -def test_azure_ai_catalog_name_prices_the_same_in_any_casing(spec: TokenPricedCatalogModel) -> None: - lowercase_cost = cost_per_token(model=f"azure_ai/{spec.catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) - upper_cost = cost_per_token(model=f"azure_ai/{spec.catalog_name.upper()}", prompt_tokens=A_MILLION, completion_tokens=0) +@pytest.mark.parametrize("catalog_name", TOKEN_PRICED_NAMES) +def test_azure_ai_catalog_name_prices_the_same_in_any_casing(catalog_name: str) -> None: + lowercase_cost = cost_per_token(model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + upper_cost = cost_per_token(model=f"azure_ai/{catalog_name.upper()}", prompt_tokens=A_MILLION, completion_tokens=0) assert upper_cost == lowercase_cost @@ -80,19 +89,10 @@ def test_azure_ai_grok_4_20_bills_cached_prompt_tokens_at_the_input_price(catalo @pytest.mark.usefixtures("local_model_cost_map") def test_azure_ai_whisper_catalog_name_is_priced_per_second() -> None: - transcription: Final = TranscriptionResponse(text="hello") - transcription._hidden_params = { # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter - "custom_llm_provider": "azure_ai", - "model": "azure_ai/whisper", - "audio_transcription_duration": 3600, - } - cost = completion_cost( - completion_response=transcription, - model="azure_ai/whisper", - custom_llm_provider="azure_ai", - call_type="atranscription", - ) - assert cost == pytest.approx(0.36) + one_second_cost: Final = _whisper_transcription_cost(1) + one_hour_cost: Final = _whisper_transcription_cost(AN_HOUR_IN_SECONDS) + assert one_second_cost > 0 + assert one_hour_cost == pytest.approx(AN_HOUR_IN_SECONDS * one_second_cost) @pytest.mark.parametrize("catalog_name", CATALOG_NAMES) From 3cadf2f8f7120bf10409a353ef08e4cdc6f78b80 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:35:19 -0700 Subject: [PATCH 026/157] test(azure_ai): charge the router fee over cached prompt tokens too --- .../llms/azure_ai/test_azure_ai_cost_calculator.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 7df14b91741..a43fc3332af 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -138,6 +138,18 @@ class TestAzureModelRouterFlatCost: assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) assert completion_cost_usd == 0.0 + def test_unmapped_router_deployment_name_charges_the_fee_over_cached_prompt_tokens_too(self) -> None: + usage = Usage( + prompt_tokens=2000, + completion_tokens=800, + total_tokens=2800, + cache_read_input_tokens=500, + cache_creation_input_tokens=200, + ) + prompt_cost, completion_cost_usd = cost_per_token(model="azure-model-router", usage=usage) + assert prompt_cost == pytest.approx(2000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 + def test_router_deployment_name_as_both_names_charges_the_fee_once(self) -> None: usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) prompt_cost, completion_cost_usd = cost_per_token( From 78893723f7b65cb73d2c5285cb412b25764a3857 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 13:05:38 +0200 Subject: [PATCH 027/157] fix(auth): drop env-credential hint from 401 when env login is disabled --- litellm/proxy/auth/login_utils.py | 7 ++++++- tests/test_litellm/proxy/auth/test_login_utils.py | 3 +++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 6ca7ea5074d..c0a76a4fc20 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -362,8 +362,13 @@ async def authenticate_user( code=401, ) else: + env_credentials_hint: Final = ( + "\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file" + if is_env_credential_login_enabled(general_settings) + else "" + ) raise ProxyException( - message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", + message=f"Invalid credentials used to access UI.{env_credentials_hint}", type=ProxyErrorTypes.auth_error, param="invalid_credentials", code=401, diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 72cc6e04a12..e209a491b0a 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -186,6 +186,7 @@ async def test_authenticate_user_invalid_credentials(): assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" assert "Invalid credentials" in exc_info.value.message + assert "UI_USERNAME" in exc_info.value.message @pytest.mark.asyncio @@ -829,6 +830,8 @@ class TestDisableEnvCredentialLogin: assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" + assert "UI_USERNAME" not in exc_info.value.message + assert "UI_PASSWORD" not in exc_info.value.message @pytest.mark.asyncio async def test_rejects_master_key_fallback_when_disabled(self): From 992ee37258e7f3c9ed13c374610899388bc8d446 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 13:21:45 +0200 Subject: [PATCH 028/157] Update litellm/proxy/health_endpoints/_health_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/health_endpoints/_health_endpoints.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e5e4f89234b..747cfa526d9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1583,16 +1583,6 @@ async def _show_no_redis_warning() -> bool: def _show_env_credential_login_warning() -> bool: - """ - Whether the UI should warn admins that env-credential login is still enabled. - - UI_USERNAME/UI_PASSWORD (or the master key, when UI_PASSWORD is unset) grant - proxy-admin access with a shared static secret: no per-person identity, no - audit trail, no password policy, and it stays valid until the env var or - master key rotates. That is fine for first-time setup, so it is on by - default, but once real admin accounts exist it should be turned off with - `general_settings.disable_env_credential_login`. - """ from litellm.proxy.auth.login_utils import is_env_credential_login_enabled from litellm.proxy.proxy_server import general_settings From cb8daee55f05e3d4ec728c1b3f15e153176721ea Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 8 Sep 2026 13:23:25 +0000 Subject: [PATCH 029/157] chore(registry): absorb #40185 daybreak-blue and #40159 cloudflare rpm limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 93 ++++++++++++++++--- model_prices_and_context_window.json | 93 ++++++++++++++++--- tests/test_litellm/test_utils.py | 6 ++ 3 files changed, 164 insertions(+), 28 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5de376218b9..ef8da798922 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13792,7 +13792,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13801,7 +13802,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13810,7 +13812,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -13819,7 +13822,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -13829,6 +13833,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13839,7 +13844,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -13848,7 +13854,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -13857,7 +13864,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -13866,7 +13874,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -13877,6 +13886,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13888,6 +13898,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -13897,7 +13908,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -13906,7 +13918,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -13917,6 +13930,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13928,6 +13942,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13938,7 +13953,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -13948,6 +13964,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -13958,6 +13975,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -13967,7 +13985,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -13978,6 +13997,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13989,6 +14009,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13999,7 +14020,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -14009,6 +14031,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14019,7 +14042,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -14029,6 +14053,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14040,6 +14065,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14050,6 +14076,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14060,6 +14087,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14071,6 +14099,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14081,6 +14110,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -54901,6 +54931,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -60822,6 +60885,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60832,6 +60896,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5de376218b9..ef8da798922 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13792,7 +13792,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13801,7 +13802,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13810,7 +13812,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -13819,7 +13822,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -13829,6 +13833,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13839,7 +13844,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -13848,7 +13854,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -13857,7 +13864,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -13866,7 +13874,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -13877,6 +13886,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13888,6 +13898,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -13897,7 +13908,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -13906,7 +13918,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -13917,6 +13930,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13928,6 +13942,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13938,7 +13953,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -13948,6 +13964,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -13958,6 +13975,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -13967,7 +13985,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -13978,6 +13997,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13989,6 +14009,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13999,7 +14020,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -14009,6 +14031,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14019,7 +14042,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -14029,6 +14053,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14040,6 +14065,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14050,6 +14076,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14060,6 +14087,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14071,6 +14099,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14081,6 +14110,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -54901,6 +54931,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -60822,6 +60885,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60832,6 +60896,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8a56a84ade7..cfdbbcb1132 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -54,6 +54,12 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_cloudflare_model_info_includes_rpm(local_model_cost_map): + assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 + assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 + assert litellm.get_model_info("cloudflare/@cf/openai/whisper-large-v3-turbo")["rpm"] == 720 + + def test_get_utc_datetime_returns_current_aware_utc_time() -> None: before: Final = datetime.now(timezone.utc) result: Final = litellm.utils.get_utc_datetime() From fd1fad5e0559db253842fdb3972a94081c58734d Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 8 Sep 2026 13:37:56 +0000 Subject: [PATCH 030/157] test: annotate registry metadata test parameters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/test_chatgpt_responses_transformation.py | 6 ++++-- tests/test_litellm/test_utils.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8a0027fdb53..a7520bd5955 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -54,7 +54,7 @@ class TestChatGPTResponsesAPITransformation: "chatgpt/gpt-5.6-terra", ], ) - def test_chatgpt_responses_model_metadata(self, model_name, local_model_cost_map): + def test_chatgpt_responses_model_metadata(self, model_name: str, local_model_cost_map: None) -> None: model_info = litellm.get_model_info(model_name) assert model_info["litellm_provider"] == "chatgpt" @@ -75,7 +75,9 @@ class TestChatGPTResponsesAPITransformation: "gpt-5.6-terra", ], ) - def test_chatgpt_models_bridge_chat_completions_to_responses(self, model_name, local_model_cost_map): + def test_chatgpt_models_bridge_chat_completions_to_responses( + self, model_name: str, local_model_cost_map: None + ) -> None: """A chat completions request for these models must take the Responses bridge. `gpt-5.6-*` also exists as an openai chat model, so an unregistered diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index cfdbbcb1132..45532f91f16 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -54,7 +54,7 @@ from litellm.utils import ( # Adds the parent directory to the system path -def test_cloudflare_model_info_includes_rpm(local_model_cost_map): +def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 assert litellm.get_model_info("cloudflare/@cf/openai/whisper-large-v3-turbo")["rpm"] == 720 From 17dd6afbaafed87e6e7380278ca13a543046160a Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 13:51:58 +0000 Subject: [PATCH 031/157] fix: keep reasoning_effort for mode: responses bridge deployments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/main.py | 12 +++++++++++- tests/test_litellm/test_main.py | 34 +++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..062ddb9a293 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5398,6 +5398,16 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params + allowed_openai_params: Final[list[str] | None] = cast( + list[str] | None, + ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if responses_api_model_info.get("mode") == "responses" + and not skip_responses_api_bridge + and "reasoning_effort" not in (kwargs.get("allowed_openai_params") or []) + else kwargs.get("allowed_openai_params") + ), + ) optional_param_args: Final = { "functions": functions, "function_call": function_call, @@ -5442,7 +5452,7 @@ def completion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, - "allowed_openai_params": kwargs.get("allowed_openai_params"), + "allowed_openai_params": allowed_openai_params, "base_model": base_model, } optional_params = get_optional_params(**optional_param_args, **non_default_params) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 038df3656fe..9ce2bb8400d 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1367,6 +1367,40 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( } +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + mock_responses_completion, reasoning_effort, restore_model_registry +): + mock_responses_completion.return_value = MagicMock() + model = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": True, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + with patch.object(litellm, "supports_reasoning", return_value=False): + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == reasoning_effort + + @pytest.mark.parametrize( "model, model_info, expected_model_param, expected_base_model_param", [ From 5520c3bb1a768263bbc4458b14da7d87e3696a41 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 13:53:02 +0000 Subject: [PATCH 032/157] refactor: simplify allowed_openai_params bridge override Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/main.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 062ddb9a293..01fe6a7c674 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5398,15 +5398,13 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params - allowed_openai_params: Final[list[str] | None] = cast( - list[str] | None, - ( - [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] - if responses_api_model_info.get("mode") == "responses" - and not skip_responses_api_bridge - and "reasoning_effort" not in (kwargs.get("allowed_openai_params") or []) - else kwargs.get("allowed_openai_params") - ), + bridges_to_responses_api: Final = ( + responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge + ) + allowed_openai_params: Final[list[str] | None] = ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if bridges_to_responses_api + else kwargs.get("allowed_openai_params") ) optional_param_args: Final = { "functions": functions, From 6f3b3c7957a67d45542538bcda6dc689752b7316 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 14:27:12 +0000 Subject: [PATCH 033/157] test: cover responses bridge reasoning at HTTP boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_main.py | 70 +++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 16 deletions(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 9ce2bb8400d..8a9d5e07258 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1368,18 +1368,57 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( @pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) -@patch("litellm.completion_extras.responses_api_bridge.completion") def test_responses_bridge_preserves_reasoning_effort_with_drop_params( - mock_responses_completion, reasoning_effort, restore_model_registry + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ): - mock_responses_completion.return_value = MagicMock() - model = "perplexity/test-responses-bridge" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" litellm.register_model( { model: { "litellm_provider": "perplexity", "mode": "responses", - "supports_reasoning": True, + "supports_reasoning": False, "input_cost_per_token": 0.0, "output_cost_per_token": 0.0, } @@ -1387,18 +1426,17 @@ def test_responses_bridge_preserves_reasoning_effort_with_drop_params( persist_across_reloads=False, ) - with patch.object(litellm, "supports_reasoning", return_value=False): - litellm.completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - reasoning_effort=reasoning_effort, - drop_params=True, - api_key="fake-key", - api_base="https://api.perplexity.ai", - ) + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == reasoning_effort + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} @pytest.mark.parametrize( From 06497ab42a1ad055e8cbb7695d778cd7a4158ea5 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 15:13:53 +0000 Subject: [PATCH 034/157] fix(ui): guard playground cost metric against null and NaN Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_ui/ResponseMetrics.test.tsx | 19 ++++++++++++ .../components/chat_ui/ResponseMetrics.tsx | 2 +- .../llm_calls/chat_completion.test.tsx | 6 ++++ .../components/llm_calls/chat_completion.tsx | 5 +++- .../llm_calls/responses_api.test.tsx | 29 +++++++++++++++++++ .../components/llm_calls/responses_api.tsx | 5 +++- 6 files changed, 63 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx index f31e1839739..4ba6d8d50b4 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx @@ -51,4 +51,23 @@ describe("ResponseMetrics prompt cache chips", () => { expect(screen.queryByText(/Response Cache/)).not.toBeInTheDocument(); }); + + it("does not render the Cost chip when a persisted cost is null", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + expect(screen.getByText("In: 1")).toBeInTheDocument(); + }); + + it("does not render the Cost chip for NaN", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + }); + + it("renders the Cost chip for a finite cost", () => { + render(); + + expect(screen.getByText("Cost: $0.000063")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx index ec62d0618d7..bb7debc7aee 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx @@ -159,7 +159,7 @@ const ResponseMetrics: React.FC = ({ timeToFirstToken, tot /> )} - {usage?.cost !== undefined && ( + {typeof usage?.cost === "number" && Number.isFinite(usage.cost) && ( { expect(usageData).not.toHaveProperty("cacheReadTokens"); expect(usageData).not.toHaveProperty("cacheCreationTokens"); }); + + it("omits cost when the provider reports a non-numeric value", async () => { + const usageData = await captureUsage({ cost: "not-a-number" }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); }); describe("chat_completion response cache", () => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index cd0852bc06e..b813a35f687 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -245,7 +245,10 @@ export async function makeOpenAIChatCompletionRequest( // Extract cost from usage object if available if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - usageData.cost = parseFloat(chunkWithUsage.usage.cost); + const parsedCost = parseFloat(chunkWithUsage.usage.cost); + if (Number.isFinite(parsedCost)) { + usageData.cost = parsedCost; + } } onUsageData(usageData); diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 0b94c093acd..7b13f0ac2e3 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -231,6 +231,35 @@ describe("responses_api", () => { expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); }); + it("should omit cost when the proxy reports a non-numeric cost", async () => { + async function* streamWithNonNumericCost() { + yield { + type: "response.completed", + response: { + id: "resp_non_numeric_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: "not-a-number" }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithNonNumericCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + it("should replay MCP output items as events for a non-streaming response", async () => { mockResponsesCreate.mockReturnValueOnce( nonStreamingResponse({ diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 7ab76488504..63085a81b94 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -312,7 +312,10 @@ export async function makeOpenAIResponsesRequest( } if (usage.cost !== undefined && usage.cost !== null) { - usageData.cost = Number(usage.cost); + const parsedCost = Number(usage.cost); + if (Number.isFinite(parsedCost)) { + usageData.cost = parsedCost; + } } onUsageData(usageData, mcpToolUsed); From cbe340a31ca81c40be088520ebbaaaca644977ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:38:18 -0700 Subject: [PATCH 035/157] feat(guardrails): deliver tool-call rewrites into buffered streams A post_call pipeline guardrail that rewrites a streamed tool call (its arguments or its name) now has that rewrite written back across the buffered chunks on chat, Responses, and Messages streams, so the client receives the rewritten tool call instead of the original. The chat handler rewrites the first fragment of each tool-call index and blanks the rest, the Responses handler syncs the function_call output items and their argument events, and the Messages handler rewrites the tool_use content_block_start and input_json_delta events in both dict and SSE-bytes chunks. The delivers_ended_stream_text_rewrites flag becomes delivers_ended_stream_rewrites, since the write-back now covers both text and tool calls, and the executor only discards a tool-call rewrite on translations without write-back or on a shape the translation refuses. --- .../chat/guardrail_translation/handler.py | 176 +++++++++++++++--- .../guardrail_translation/base_translation.py | 17 +- .../chat/guardrail_translation/handler.py | 96 +++++++++- .../guardrail_translation/handler.py | 111 ++++++++++- .../proxy/policy_engine/pipeline_executor.py | 22 +-- litellm/proxy/utils.py | 10 +- .../test_anthropic_guardrail_handler.py | 66 +++++++ .../test_openai_guardrail_handler.py | 73 ++++++++ ...test_openai_responses_guardrail_handler.py | 88 +++++++++ .../policy_engine/test_pipeline_executor.py | 22 ++- .../proxy_logging/test_guardrail_pipeline.py | 9 +- 11 files changed, 618 insertions(+), 72 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e486be12fe2..2049866b444 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,10 +13,11 @@ Pattern Overview: """ import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_field, stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -153,6 +155,28 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +@dataclass(frozen=True, slots=True) +class _ToolCallShape: + name: str | None + arguments: str + + +_SSEEventRewriter = Callable[[Mapping[str, object]], Mapping[str, object] | None] + + +def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: + """The guardrail-visible shape of each tool call, whether the guardrail handed + back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" + functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls) + return tuple( + _ToolCallShape( + name=name if isinstance(name := stream_item_field(function, "name"), str) else None, + arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "", + ) + for function in functions + ) + + class _AnthropicSSEDelta(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] @@ -170,7 +194,7 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def __init__(self): super().__init__() @@ -1050,6 +1074,7 @@ class AnthropicMessagesHandler(BaseTranslation): first_choice.message.tool_calls, ) string_so_far = first_choice.message.content + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() if string_so_far: guardrail_inputs["texts"] = [string_so_far] @@ -1084,6 +1109,19 @@ class AnthropicMessagesHandler(BaseTranslation): and guardrailed_texts[0] != string_so_far ): self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) + if deliver_ended_stream_rewrites: + returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls") + self._write_ended_stream_tool_call_rewrites( + responses_so_far, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=_tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tool_calls_list or () + ), + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1212,38 +1250,120 @@ class AnthropicMessagesHandler(BaseTranslation): """Deliver an ended-stream guardrail text rewrite by rewriting the buffered chunks in place: the first ``text_delta`` carries the full rewritten text and every later one is blanked, leaving the surrounding - message and content-block framing untouched. Handles both chunk formats - this stream carries (parsed event dicts and raw SSE bytes).""" + message and content-block framing untouched.""" replacements: Final = chain((rewritten_text,), repeat("")) - for idx, item in enumerate(responses_so_far): - if isinstance(item, dict): - delta = item.get("delta") - if item.get("type") == "content_block_delta" and isinstance(delta, dict): - if delta.get("type") == "text_delta": - delta["text"] = next(replacements) - elif isinstance(item, (bytes, bytearray)): - responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer - AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements) - ) + + def rewrite_text_delta(event: Mapping[str, object]) -> Mapping[str, object] | None: + delta: Final = event.get("delta") + if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): + return None + if delta.get("type") != "text_delta": + return None + return {**event, "delta": {**delta, "text": next(replacements)}} + + AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + + @classmethod + def _write_ended_stream_tool_call_rewrites( + cls, + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + *, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Deliver ended-stream guardrail tool-call rewrites by rewriting the + buffered chunks in place: the rebuilt response lists tool calls in the + order of the stream's ``tool_use`` blocks, so the nth rewritten call lands + on the nth block, its first ``input_json_delta`` carrying the full rewritten + arguments, every later one blanked, and ``content_block_start`` carrying the + rewritten name. Blocks that do not line up with the rebuilt tool calls make + the rewrite undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + block_indices: Final = tuple( + index + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + and isinstance(index := event.get("index"), int) + ) + if len(block_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_block: Final = MappingProxyType( + { + index: after + for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + argument_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} + ) + + def rewrite_tool_use(event: Mapping[str, object]) -> Mapping[str, object] | None: + index: Final = event.get("index") + if not isinstance(index, int) or index not in rewrites_by_block: + return None + match event.get("type"): + case "content_block_start": + block: Final = event.get("content_block") + name: Final = rewrites_by_block[index].name + if not isinstance(block, Mapping) or name is None: + return None + return {**event, "content_block": {**block, "name": name}} + case "content_block_delta": + delta: Final = event.get("delta") + if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": + return None + return {**event, "delta": {**delta, "partial_json": next(argument_replacements[index])}} + case _: + return None + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) @staticmethod - def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes: - """Rewrite every ``text_delta`` data line in one SSE chunk with the next - replacement text, leaving all other events and framing byte-identical.""" + def _rewrite_ended_stream_events( + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + rewrite_event: _SSEEventRewriter, + ) -> None: + """Replace every buffered event ``rewrite_event`` returns a rewrite for, in + both chunk formats this stream carries (parsed event dicts and raw SSE + bytes), leaving every other event and the framing untouched.""" + rewritten_items: Final = tuple( + AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far + ) + responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer + + @staticmethod + def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: + if isinstance(item, dict): + rewritten: Final = rewrite_event(_as_str_mapping(item)) + return item if rewritten is None else dict(rewritten) + if isinstance(item, (bytes, bytearray)): + return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) + return item + + @staticmethod + def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes: + """Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites, + leaving all other events and framing byte-identical.""" try: decoded: Final = sse_bytes.decode("utf-8") except UnicodeDecodeError: return sse_bytes return "\n\n".join( - AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n") + "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n")) + for block in decoded.split("\n\n") ).encode("utf-8") @staticmethod - def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str: - return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n")) - - @staticmethod - def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str: + def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str: if not line.startswith("data:"): return line try: @@ -1252,14 +1372,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) except json.JSONDecodeError: return line - if not isinstance(data, dict) or data.get("type") != "content_block_delta": + if not isinstance(data, dict): return line - delta: Final = data.get("delta") - if not isinstance(delta, dict) or delta.get("type") != "text_delta": - return line - return "data: " + json.dumps( - {**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts - ) + rewritten: Final = rewrite_event(_as_str_mapping(data)) + return line if rewritten is None else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index afd8e0f67f7..6d1a9ab1c3e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -52,13 +52,14 @@ class StreamingScanKey: class BaseTranslation(ABC): - delivers_ended_stream_text_rewrites: ClassVar[bool] = False + delivers_ended_stream_rewrites: ClassVar[bool] = False """Whether ``process_output_streaming_response`` accepts ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) - stream, writes guardrail text rewrites back across ``responses_so_far`` so - a buffered pipeline can release rewritten chunks. Tool-call rewrites, and - text rewrites on every other translation, are undeliverable: the pipeline - executor discards them and releases the original chunks.""" + stream, writes guardrail text and tool-call rewrites back across + ``responses_so_far`` so a buffered pipeline can release rewritten chunks, + raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites + on every other translation are undeliverable: the pipeline executor + discards them and releases the original chunks.""" @staticmethod def transform_user_api_key_dict_to_metadata( @@ -175,9 +176,9 @@ class BaseTranslation(ABC): transformations (see ``StreamTransformSink``); base handlers ignore it. ``deliver_ended_stream_rewrites`` is passed True only when the caller holds the whole buffered stream and the subclass declares - ``delivers_ended_stream_text_rewrites``: the handler then writes - guardrail text rewrites back across ``responses_so_far`` instead of - discarding them. + ``delivers_ended_stream_rewrites``: the handler then writes + guardrail text and tool-call rewrites back across ``responses_so_far`` + instead of discarding them. """ return responses_so_far diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 80292aef2cf..42c95ac3316 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import coerce_stream_holdback_value, ) from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, Choices, GenericGuardrailAPIInputs, ModelResponse, @@ -78,7 +80,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -610,13 +612,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): deliver_ended_stream_rewrites: bool, ) -> None: """Ended-stream path: rebuild the full response, run the non-streaming - output guardrail against it, and (when opted in) write any text rewrite - back across the buffered chunks.""" + output guardrail against it, and (when opted in) write any text or + tool-call rewrite back across the buffered chunks.""" model_response: Final = cast( ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), ) pre_guardrail_texts: Final = self._string_choice_contents(model_response) + pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) await self.process_output_response( response=model_response, guardrail_to_apply=guardrail_to_apply, @@ -624,13 +627,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict=user_api_key_dict, request_data=request_data, ) - if deliver_ended_stream_rewrites: - await self._write_ended_stream_text_rewrites( - responses_so_far=responses_so_far, - guardrailed_response=model_response, - pre_guardrail_texts=pre_guardrail_texts, - guardrail_name=guardrail_to_apply.guardrail_name or "unknown", - ) + if not deliver_ended_stream_rewrites: + return + guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown" + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_name, + ) + self._write_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_name, + ) def build_stream_error_items( self, @@ -1043,6 +1054,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists ) + @staticmethod + def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]: + return tuple( + (tool_call.function.name, tool_call.function.arguments) + for choice in response.choices + for tool_call in choice.message.tool_calls or () + if isinstance(tool_call, ChatCompletionMessageToolCall) + ) + + @staticmethod + def _function_tool_call_fragments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" + fragments: Final = tuple( + tool_call + for response in responses_so_far + for choice in response.choices + for tool_call in choice.delta.tool_calls or () + if isinstance(tool_call, ChatCompletionDeltaToolCall) + ) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + return tuple( + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + ) + + def _write_ended_stream_tool_call_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for before, (name, arguments), fragments in zip( + pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call + ): + if (name, arguments) == before: + continue + head, *tail = fragments + head.function.name = name + head.function.arguments = arguments + for fragment in tail: + fragment.function.arguments = "" + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b0f79552bc5..447870fc0be 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -101,6 +101,18 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseInputParam +class _ToolCallShape(NamedTuple): + name: str | None + arguments: str + + +def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: + return tuple( + _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) + for tool_call in tool_calls + ) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -340,7 +352,7 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -754,6 +766,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -762,6 +775,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) + else tool_calls_to_check + ) # Write guardrailed texts back into the output items in-place. # final_chunk is a reference into responses_so_far so this @@ -784,6 +803,13 @@ class OpenAIResponsesHandler(BaseTranslation): stream_events=responses_so_far[:-1], rewrites_by_position=rewrites_by_position, ) + self._deliver_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + outputs=outputs, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) return responses_so_far # ------------------------------------------------------------------ # @@ -894,6 +920,89 @@ class OpenAIResponsesHandler(BaseTranslation): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + def _deliver_ended_stream_tool_call_rewrites( + self, + responses_so_far: Sequence[object], + outputs: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites into the completed + envelope's ``function_call`` items and sync the earlier stream events, + keyed by ``output_index``. The guardrail sees the envelope's function + calls in output order, which is how a rewritten call finds its item; a + rewrite whose calls do not line up with the envelope is reported as + undeliverable, so the pipeline executor discards it and releases the + original events.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + function_call_indices: Final = tuple( + output_idx + for output_idx, output_item in enumerate(outputs) + if stream_item_field(output_item, "type") == "function_call" + ) + if len(function_call_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_output_index: Final = MappingProxyType( + { + output_idx: after + for output_idx, before, after in zip( + function_call_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls + ) + if after != before + } + ) + for output_idx, rewrite in rewrites_by_output_index.items(): + self._write_function_call_item(outputs[output_idx], rewrite.name, rewrite.arguments) + self._sync_stream_events_with_tool_call_rewrites( + stream_events=responses_so_far[:-1], + rewrites_by_output_index=rewrites_by_output_index, + ) + + def _sync_stream_events_with_tool_call_rewrites( + self, + stream_events: Sequence[object], + rewrites_by_output_index: Mapping[int, _ToolCallShape], + ) -> None: + """Sync pre-completion function-call events with the rewritten completed + response: the first ``function_call_arguments.delta`` for a rewritten call + carries the full rewritten arguments and the rest are blanked, while + ``function_call_arguments.done`` and ``output_item.done`` carry the full + rewritten arguments and ``output_item.added`` / ``output_item.done`` the + rewritten name, so every event a client may read agrees with the + rewritten ``response.completed`` payload.""" + delta_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_output_index.items()} + ) + for event in stream_events: + output_index = stream_item_field(event, "output_index") + if not isinstance(output_index, int) or output_index not in rewrites_by_output_index: + continue + rewrite = rewrites_by_output_index[output_index] + match stream_item_field(event, "type"): + case "response.function_call_arguments.delta": + self._write_event_field(event, "delta", next(delta_replacements[output_index])) + case "response.function_call_arguments.done": + self._write_event_field(event, "arguments", rewrite.arguments) + case "response.output_item.added": + self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, None) + case "response.output_item.done": + self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, rewrite.arguments) + case _: + pass + + @staticmethod + def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: + if not (isinstance(item, dict) or hasattr(item, "get")): + return + if name is not None: + OpenAIResponsesHandler._write_event_field(item, "name", name) + if arguments is not None: + OpenAIResponsesHandler._write_event_field(item, "arguments", arguments) + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9bc10949e9f..3e112bf8e67 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -85,10 +85,10 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT: class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's guardrail. It records whether the guardrail returned different output than it was given, - which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text - rewrites are deliverable on translations that write them back across the buffered chunks - (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any - other translation are discarded by the executor, which releases the original chunks. + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and + tool-call rewrites are deliverable on translations that write them back across the + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation + are discarded by the executor, which releases the original chunks. The inner guardrail's ``apply_guardrail`` already records the guardrail information and span, so the observer's stays out of ``log_guardrail_information``.""" @@ -290,13 +290,13 @@ class PipelineExecutor: litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back. A rewrite that - cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation - without write-back, or one the translation refused with - ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the - originals and the step passes, so the client gets the stream the merge base sent.""" + text and tool-call rewrites on translations that support ended-stream write-back. A + rewrite that cannot reach the client yet (one on a translation without write-back, or + one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the + buffered chunks go back to the originals and the step passes, so the client gets the + stream the merge base sent.""" observer: Final = _StreamRewriteObserver(callback) - deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites originals: Final = copy.deepcopy(streaming_chunks) try: if deliver_rewrites: @@ -319,7 +319,7 @@ class PipelineExecutor: except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) else: - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + if not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls): _release_original_chunks(step.guardrail, streaming_chunks, originals) if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..76b9b4b44ed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3445,11 +3445,11 @@ class ProxyLogging: assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow releases the buffered chunks: verbatim when no guardrail rewrote the - output, rewritten in place when one rewrote text and the translation - delivers ended-stream rewrites (later steps then re-scan the rewritten - chunks, so rewrites chain). A rewrite the translation cannot deliver - yet (a tool-call rewrite, or a text rewrite on a route without - write-back) is discarded by the executor and the original chunks are + output, rewritten in place when one rewrote text or a tool call and the + translation delivers ended-stream rewrites (later steps then re-scan the + rewritten chunks, so rewrites chain). A rewrite the translation cannot + deliver yet (one on a route without write-back, or a shape the route + refuses) is discarded by the executor and the original chunks are released, as is a buffered shape no translation resolves; a block or modify_response terminates with the translation's block chunks or the raised error. 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 bf40f781fa3..bd2dff9b33c 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 @@ -315,6 +315,72 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert "event: message_start" in raw and "event: message_stop" in raw assert '"stop_reason": "end_turn"' in raw + @staticmethod + def _ended_tool_use_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + return MaskArguments(guardrail_name="test") + + @staticmethod + def _partial_jsons(chunks: list) -> list: + return [ + json.loads(line[len("data:") :].strip())["delta"]["partial_json"] + for chunk in chunks + for line in chunk.decode().split("\n") + if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""] + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert "persim" not in raw + + @pytest.mark.asyncio + async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index aff0530ee8b..be168ec83e6 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1113,6 +1113,79 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[1].choices[0].delta.content in (None, "") assert chunks[1].choices[0].finish_reason == "stop" + @staticmethod + def _ended_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None): + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(fragment("", name="lookup_fruit", call_id="call_1")), + chunk(fragment('{"fruit":')), + chunk(fragment(' "persimmon"}')), + chunk(None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""] + assert fragments[0][0].function.name == "lookup_fruit" + assert fragments[0][0].id == "call_1" + assert chunks[3].choices[0].delta.tool_calls is None + assert chunks[3].choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}'] + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 33c8c97fea7..a024fc7ff81 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1195,6 +1195,94 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @staticmethod + def _ended_function_call_stream_events() -> List[dict]: + def item(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_123", + "call_id": "call_123", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_123", + "output_index": 0, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [item('{"fruit": "persimmon"}', "completed")], + "status": "completed", + }, + }, + ] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + return MaskArguments(guardrail_name="test-mask-arguments") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["arguments"] == "" + assert events[1]["delta"] == '{"fruit": "[MASKED]"}' + assert events[2]["delta"] == "" + assert events[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + + @pytest.mark.asyncio + async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[1]["delta"] == '{"fruit":' + assert events[3]["arguments"] == '{"fruit": "persimmon"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}' + @pytest.mark.asyncio @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 54ef1f79f4d..ee75003db9f 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -924,7 +924,7 @@ class _TextReturningGuardrail(CustomGuardrail): class _TextTranslation: - delivers_ended_stream_text_rewrites = False + delivers_ended_stream_rewrites = False def __init__(self): self.seen_guardrail_names = [] @@ -946,7 +946,7 @@ class _WritingTranslation: """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the chat/Responses/Messages handlers do on an ended stream.""" - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -970,7 +970,7 @@ class _WritingTranslation: class _RefusingTranslation: - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -1088,13 +1088,27 @@ async def test_streaming_step_delivers_text_rewrite_through_writing_translation( @pytest.mark.asyncio -async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog): +async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) + _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 73dd6746b29..d27f504ba6d 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1836,7 +1836,7 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite( +async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog ): transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 @@ -1855,9 +1855,12 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_to delivered.append(item) assert len(delivered) == 2 - assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0] + assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}' + assert delivered_tool_call.function.name == "lookup" + assert delivered_tool_call.id == "call_1" assert delivered[1].choices[0].finish_reason == "tool_calls" - assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + assert not any("discarded" in message for message in _warnings(caplog)) @pytest.mark.asyncio From e65c56876f9aae1f94120212a2092592792f6538 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:42:51 -0700 Subject: [PATCH 036/157] fix(guardrails): write tool-call rewrites into typed Responses envelopes The response.completed envelope carries its function_call output items as SDK objects without a get shim, so the write-back skipped them and the envelope still showed the original arguments after every stream event had been rewritten. Write the item whenever one is present, and cover the typed event shape the live proxy carries in the handler test. --- .../guardrail_translation/handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 46 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 447870fc0be..208ceb959e6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -996,7 +996,7 @@ class OpenAIResponsesHandler(BaseTranslation): @staticmethod def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: - if not (isinstance(item, dict) or hasattr(item, "get")): + if item is None: return if name is not None: OpenAIResponsesHandler._write_event_field(item, "name", name) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a024fc7ff81..bdf7f83757c 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1222,6 +1222,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: "type": "response.completed", "response": { "id": "resp_123", + "created_at": 1, "model": "gpt-4o", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed", @@ -1268,6 +1269,51 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self): + from litellm.types.llms.openai import ( + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[Any] = [ + model.model_validate(event) + for model, event in zip( + ( + OutputItemAddedEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_function_call_stream_events(), + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event, ResponseCompletedEvent) + assert isinstance(completed_event.response, ResponsesAPIResponse) + assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == '{"fruit": "[MASKED]"}' + assert typed_events[2].delta == "" + assert typed_events[3].arguments == '{"fruit": "[MASKED]"}' + assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].name == "lookup_fruit" + @pytest.mark.asyncio async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): handler = OpenAIResponsesHandler() From 3caa3b60d5a319efa78871198081783dd4c70ad1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:58:18 -0700 Subject: [PATCH 037/157] feat(guardrails): run post_call policy pipelines on background Responses retrieval A POST /v1/responses with background: true returns a queued response, so the post_call pipelines attached at submit time had nothing to inspect. They now defer on queued and in_progress responses and run on GET /v1/responses/{id} instead: the retrieval resolves the response id back to its deployment, re-attaches the policies that governed the original model, and reports them in the x-litellm-applied-* headers of the retrieval response. --- litellm/proxy/common_request_processing.py | 8 +- .../proxy/policy_engine/response_retrieval.py | 116 +++++++++++++ litellm/proxy/utils.py | 43 +++-- .../policy_engine/test_response_retrieval.py | 157 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 109 ++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 95 +++++++---- 6 files changed, 485 insertions(+), 43 deletions(-) create mode 100644 litellm/proxy/policy_engine/response_retrieval.py create mode 100644 tests/test_litellm/proxy/policy_engine/test_response_retrieval.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9720e4b1cf8..1c1eda73c0c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -184,6 +184,7 @@ from litellm.proxy.litellm_pre_call_utils import ( refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, ) +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -1883,7 +1884,6 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, ) - # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request. Ends at start_time # (not a freshly captured time.time() here) so this window is exactly @@ -2031,6 +2031,12 @@ class ProxyBaseLLMRequestProcessing: data=self.data, call_type=route_type, ) + if route_type == "aget_responses": + attach_post_call_pipelines_to_retrieval( + data=self.data, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py new file mode 100644 index 00000000000..54b9acf6150 --- /dev/null +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -0,0 +1,116 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, +) +from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] + +_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | None: + if llm_router is None or not isinstance(response_id, str): + return None + model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) + if model_id is None: + return None + deployment: Final = llm_router.get_deployment(model_id) + return deployment.model_name if deployment is not None else None + + +def _retrieval_context( + data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str +) -> PolicyMatchContext: + team_alias: Final = user_api_key_dict.team_alias + key_alias: Final = user_api_key_dict.key_alias + return PolicyMatchContext( + team_alias=team_alias if isinstance(team_alias, str) else None, + key_alias=key_alias if isinstance(key_alias, str) else None, + model=model_group, + tags=get_tags_from_request_body(data) or None, + ) + + +def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: + matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + if not matches: + return (), MappingProxyType({}) + applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( + policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list + context=context, + ) + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( + context=context, policy_names=applied_policy_names + ) + if pipeline.mode == "post_call" + ) + return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) + + +def attach_post_call_pipelines_to_retrieval( + data: dict, # mutable-ok: the proxy's request-state dict, written in place like every other policy engine hook + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", +) -> None: + if not get_policy_registry().is_initialized(): + return + model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) + if model_group is None: + return + context: Final = _retrieval_context(data, user_api_key_dict, model_group) + post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) + _, bucket = get_or_create_metadata_bucket(data) + already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ()) + attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached) + added: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if policy_name not in attached_policy_names + ) + if not added: + return + pipelines: Final = (*already_attached, *added) + bucket["_guardrail_pipelines"] = pipelines + bucket["_pipeline_managed_guardrails"] = frozenset( + step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps + ) + for policy_name, _pipeline in added: + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) + for _policy_name, pipeline in added: + for step in pipeline.steps: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) + add_policy_sources_to_metadata( + request_data=data, + policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict + policy_name: policy_sources[policy_name] for policy_name, _pipeline in added + }, + ) + verbose_proxy_logger.debug( + "Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s", + data.get("response_id"), + model_group, + ", ".join(policy_name for policy_name, _pipeline in added), + ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..3d03d3858ee 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -171,6 +171,7 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -525,15 +526,21 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr ) -def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("background") is not True: - return +_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress")) + + +def _is_pending_background_response(response: LLMResponseTypes) -> bool: + return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES + + +def _log_deferred_post_call_pipelines(data: Mapping[str, object], response: ResponsesAPIResponse) -> None: policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) if not policy_names: return - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines do not run on background responses yet; " - "the response is released ungoverned by them: %s", + verbose_proxy_logger.debug( + "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", + response.id, + response.status, ", ".join(policy_names), ) @@ -1956,8 +1963,6 @@ class ProxyLogging: ) try: - _warn_background_skips_post_call_pipelines(data) - # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, @@ -2927,6 +2932,24 @@ class ProxyLogging: daemon=True, ).start() + async def _run_post_call_pipelines( + self, + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + user_api_key_dict: UserAPIKeyAuth, + response: LLMResponseTypes, + ) -> LLMResponseTypes | None: + if _is_pending_background_response(response): + _log_deferred_post_call_pipelines(data, response) + return None + _, pipeline_response = await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + return pipeline_response + async def post_call_success_hook( self, data: dict, @@ -2946,11 +2969,9 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - _, pipeline_response = await self._maybe_execute_pipelines( + pipeline_response: Final = await self._run_post_call_pipelines( data=data, user_api_key_dict=user_api_key_dict, - call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", - event_hook="post_call", response=response, ) if pipeline_response is not None: diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py new file mode 100644 index 00000000000..f7f29d6b46e --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -0,0 +1,157 @@ +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.router import Deployment, LiteLLM_Params + +GOVERNED_MODEL_GROUP = "gpt-5.4-mini" +GOVERNED_MODEL_ID = "deployment-governed" +UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" +UNGOVERNED_MODEL_ID = "deployment-ungoverned" + + +class FakeRouter: + def __init__(self, deployments: dict[str, Deployment]): + self._deployments = deployments + + def get_deployment(self, model_id: str) -> Deployment | None: + return self._deployments.get(model_id) + + +def _deployment(model_group: str, model_id: str) -> Deployment: + return Deployment( + model_name=model_group, + litellm_params=LiteLLM_Params(model=f"openai/{model_group}"), + model_info={"id": model_id}, + ) + + +def _router() -> FakeRouter: + return FakeRouter( + { + GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), + UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), + } + ) + + +def _encoded_response_id(model_id: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream" + ) + + +def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]}, + } + + +@pytest.fixture +def policy_engine(): + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry.load_policies( + { + "response-governance": _pipeline_policy("output-word-filter"), + "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), + "team-governance": _pipeline_policy("team-word-filter"), + } + ) + attachment_registry.load_attachments( + [ + {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "team-governance", "teams": ["governed-team"]}, + ] + ) + yield + policy_registry.clear() + attachment_registry.clear() + + +def _retrieval_data(model_id: str) -> dict: + return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} + + +def _attached_pipelines(data: dict) -> tuple[tuple[str, str], ...]: + return tuple( + (policy_name, ",".join(step.guardrail for step in pipeline.steps)) + for policy_name, pipeline in data["litellm_metadata"]["_guardrail_pipelines"] + ) + + +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"}) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"] + assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"} + assert "model" not in data + assert "guardrails" not in data["litellm_metadata"] + + +def test_key_and_team_context_also_governs_retrieval(policy_engine): + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval( + data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router() + ) + + assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) + + +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(UNGOVERNED_MODEL_ID) + + +def test_already_attached_policy_is_not_attached_twice(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router() + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + + +@pytest.mark.parametrize( + "response_id", + ["resp_plain_upstream_id", _encoded_response_id("deployment-missing-from-router"), None], +) +def test_unresolvable_response_id_attaches_nothing(policy_engine, response_id): + data = {"response_id": response_id, "litellm_metadata": {}} + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == {"response_id": response_id, "litellm_metadata": {}} + + +def test_without_a_router_attaches_nothing(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + + +def test_without_policy_engine_attaches_nothing(): + get_policy_registry().clear() + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6acd9d7258e..9283c1ed865 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8259,3 +8259,112 @@ class TestPassthroughHeadersAcceptImmutableMappings: assert merged["content-type"] == "text/event-stream" # the excluded hop-by-hop header is still dropped assert "transfer-encoding" not in merged + + +class TestBackgroundResponseRetrievalGovernance: + """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" + + GOVERNED_MODEL_GROUP = "gpt-5.4-mini" + GOVERNED_MODEL_ID = "deployment-governed" + + @pytest.fixture + def policy_engine(self): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + get_policy_registry().load_policies( + { + "response-governance": { + "guardrails": {"add": ["output-word-filter"]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}], + }, + } + } + ) + get_attachment_registry().load_attachments( + [{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}] + ) + yield + get_policy_registry().clear() + get_attachment_registry().clear() + + def _router(self) -> MagicMock: + from litellm.types.router import Deployment, LiteLLM_Params + + router = MagicMock() + router.get_deployment.side_effect = lambda model_id: ( + Deployment( + model_name=self.GOVERNED_MODEL_GROUP, + litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"), + model_info={"id": model_id}, + ) + if model_id == self.GOVERNED_MODEL_ID + else None + ) + return router + + async def _pre_call(self, route_type: str, monkeypatch) -> dict: + from litellm.responses.utils import ResponsesAPIRequestUtils + + client_facing_response_id = "resp_opaque-client-facing-id" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream" + ) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"response_id": client_facing_response_id, "litellm_metadata": {}} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def passthrough_add_litellm_data_to_request(data, **kwargs): + return data + + async def decrypting_pre_call_hook(user_api_key_dict, data, call_type): + if data.get("response_id") == client_facing_response_id: + data["response_id"] = encoded_response_id + return data + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + passthrough_add_litellm_data_to_request, + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook) + proxy_config = MagicMock(spec=ProxyConfig) + proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type=route_type, + llm_router=self._router(), + ) + return returned_data + + @pytest.mark.asyncio + async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( + self, policy_engine, monkeypatch + ): + data = await self._pre_call("aget_responses", monkeypatch) + + assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") + pipelines = data["litellm_metadata"]["_guardrail_pipelines"] + assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [ + ("response-governance", ["output-word-filter"]) + ] + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["model"] is None + + @pytest.mark.asyncio + async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( + self, policy_engine, monkeypatch + ): + data = await self._pre_call("aresponses", monkeypatch) + + assert "_guardrail_pipelines" not in data["litellm_metadata"] + assert "applied_policies" not in data["litellm_metadata"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 73dd6746b29..7d2c232ad97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -30,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_g from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -1419,8 +1420,69 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog)) +def _background_response(status: str, text: str = "") -> ResponsesAPIResponse: + output = ( + [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else [] + ) + return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status) + + +def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]: + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + return [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + + @pytest.mark.asyncio -async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( +@pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) +async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog, pending_status +): + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(background=True) + response = _background_response(pending_status) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert "response" not in seen + assert not _warnings(caplog) + assert any( + "response-governance" in record.getMessage() and pending_status in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["completed", "incomplete"]) +async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( + proxy_logging, make_user_api_key_auth, monkeypatch, final_status +): + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = _background_response(final_status, text="kumquat") + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +@pytest.mark.asyncio +async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): monkeypatch.setattr(litellm, "callbacks", []) @@ -1436,36 +1498,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( assert out is not None assert out.get("background") is True - assert any("response-governance" in message and "background" in message for message in _warnings(caplog)) - - -@pytest.mark.asyncio -async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog -): - seen: Dict[str, Any] = {} - monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) - data = { - "model": "m", - "messages": [{"role": "user", "content": "hi"}], - "background": True, - "metadata": { - "_guardrail_pipelines": [("request-governance", pre_call)], - "_pipeline_managed_guardrails": {"gr-post"}, - }, - } - - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="aresponses", - guardrails_only=True, - ) - - assert out is not None - assert not any("background" in message for message in _warnings(caplog)) + assert not _warnings(caplog) # --------------------------------------------------------------------------- From c6f5763443504501e88113e783d2e1ec88b9264d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:04:52 -0700 Subject: [PATCH 038/157] feat(guardrails): run legacy post-call hooks as streaming pipeline steps A post_call pipeline step whose guardrail only implements the older async_post_call_success_hook used to skip the stream entirely: PR #38721 fails that shape open with a warning. The streaming step now assembles the buffered stream into the response the hook expects, runs the hook, ends the stream with the hook's exception when it raises, and delivers the hook's rewrite through the same event write-back the unified guardrails use on chat, Responses, and Messages streams (Messages gets the Anthropic shape). A stream a pipeline manages no longer runs the same hook again after the stream ends. A guardrail with neither the unified interface nor a post-call hook keeps the fail-open, as does a rewrite the buffer cannot be patched with. --- .../chat/guardrail_translation/handler.py | 5 + .../guardrail_translation/base_translation.py | 7 + litellm/proxy/common_request_processing.py | 15 +- .../proxy/policy_engine/pipeline_executor.py | 113 ++++++++- litellm/proxy/utils.py | 30 +-- .../test_anthropic_guardrail_handler.py | 26 +++ .../policy_engine/test_pipeline_executor.py | 150 +++++++++++- .../test_proxy_logging_hook_detection.py | 26 +++ .../proxy_logging/test_guardrail_pipeline.py | 220 +++++++++++++----- 9 files changed, 505 insertions(+), 87 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e486be12fe2..f61cfe58e80 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -176,6 +176,11 @@ class AnthropicMessagesHandler(BaseTranslation): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def post_call_hook_response(self, response: object) -> object: + if not isinstance(response, ModelResponse): + return response + return self.adapter.translate_openai_response_to_anthropic(response) + @staticmethod def _build_streaming_usage_response( responses_so_far: Sequence[object], diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index afd8e0f67f7..bef472b2882 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -60,6 +60,13 @@ class BaseTranslation(ABC): text rewrites on every other translation, are undeliverable: the pipeline executor discards them and releases the original chunks.""" + def post_call_hook_response(self, response: object) -> object: + """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from + the object the translation stores under ``request_data["response"]`` while scanning an + ended stream. Chat and Responses scan that shape already; a translation that scans a + different one (Messages scans an OpenAI-shaped ModelResponse) overrides this.""" + return response + @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9720e4b1cf8..7e85b66682c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3328,9 +3328,10 @@ class ProxyBaseLLMRequestProcessing: has completed. Guardrails routed through unified_guardrail are skipped, since they already ran - via its streaming iterator. Guardrails that override - async_post_call_success_hook directly run here, including those that implement - apply_guardrail but keep their native lifecycle hooks. + via its streaming iterator, and so are guardrails a post_call policy pipeline + manages, since the pipeline ran them against the buffered stream. Guardrails + that override async_post_call_success_hook directly run here, including those + that implement apply_guardrail but keep their native lifecycle hooks. This is audit-only — content has already been delivered to the client. @@ -3340,12 +3341,18 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import _check_and_merge_model_level_guardrails + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + pipeline_managed_guardrail_names, + ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) + pipeline_managed: Final = pipeline_managed_guardrail_names(captured_data, "post_call") for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue + if cb.guardrail_name in pipeline_managed: + continue if not cb.should_run_guardrail( data=guardrail_data, event_type=GuardrailEventHooks.post_call, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9bc10949e9f..25c54d72f25 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -121,6 +121,80 @@ class _StreamRewriteObserver(CustomGuardrail): return outputs +class _ScannedTextRecorder(CustomGuardrail): + def __init__(self, guardrail_name: str) -> None: + super().__init__(guardrail_name=guardrail_name) + self.texts: tuple[str, ...] | None = None + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + self.texts = _text_snapshot(inputs.get("texts")) + return inputs + + +class _LegacyHookStreamAdapter(CustomGuardrail): + """Runs a guardrail that only implements the legacy post-call hook (no unified + ``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The + endpoint translation hands it the texts it scanned plus the assembled response under + ``request_data["response"]``; the hook gets that response in the shape its route gives + non-streaming hooks, an exception it raises ends the stream through the executor's + fail/error classification, and a replacement response is re-scanned by the same translation + so its texts reach the client through the translation's ended-stream write-back. A + replacement whose scanned texts do not line up with the originals is undeliverable, so the + executor releases the original chunks.""" + + def __init__( + self, + inner: CustomGuardrail, + endpoint_translation: "BaseTranslation", + user_api_key_dict: "UserAPIKeyAuth", + ) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.endpoint_translation: Final = endpoint_translation + self.user_api_key_dict: Final = user_api_key_dict + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + replacement: Final = await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=self.endpoint_translation.post_call_hook_response(request_data.get("response")), + ) + if replacement is None: + return inputs + scanned: Final = _text_snapshot(inputs.get("texts")) + rewritten: Final = await self._scanned_texts(replacement, logging_obj) + if scanned is None or rewritten is None or len(rewritten) != len(scanned): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + return {**inputs, "texts": list(rewritten)} + + async def _scanned_texts(self, response: object, logging_obj: "LiteLLMLoggingObj | None") -> tuple[str, ...] | None: + recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") + await self.endpoint_translation.process_output_response( + response=response, + guardrail_to_apply=recorder, + litellm_logging_obj=logging_obj, + user_api_key_dict=self.user_api_key_dict, + ) + return recorder.texts + + def _prepare_hook_input( step: PipelineStep, callback: CustomGuardrail, @@ -286,16 +360,23 @@ class PipelineExecutor: endpoint_translation: "BaseTranslation", streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place hook_input: dict[str, object], # mutable-ok: same request-payload shape as data - user_api_key_dict: "UserAPIKeyAuth | None", + user_api_key_dict: "UserAPIKeyAuth", litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back. A rewrite that - cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation - without write-back, or one the translation refused with - ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the - originals and the step passes, so the client gets the stream the merge base sent.""" - observer: Final = _StreamRewriteObserver(callback) + text rewrites on translations that support ended-stream write-back. A guardrail + without the unified interface runs its legacy post-call hook against the assembled + response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the client + yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one + the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: + the buffered chunks go back to the originals and the step passes, so the client gets + the stream the merge base sent.""" + scanner: Final = ( + callback + if PipelineExecutor.supports_unified_execution(callback) + else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict) + ) + observer: Final = _StreamRewriteObserver(scanner) deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites originals: Final = copy.deepcopy(streaming_chunks) try: @@ -379,11 +460,11 @@ class PipelineExecutor: if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) elif mode == "post_call" and streaming_chunks is not None: - if not use_unified or endpoint_translation is None: + if endpoint_translation is None: return ( "error", None, - f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation", None, ) await PipelineExecutor._run_streaming_step( @@ -433,10 +514,20 @@ class PipelineExecutor: @staticmethod def supports_unified_execution(callback: CustomGuardrail) -> bool: - """Whether this guardrail runs through the unified apply_guardrail path, - the interface streaming pipeline execution requires.""" + """Whether this guardrail runs through the unified apply_guardrail path.""" return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @staticmethod + def supports_streaming_execution(callback: CustomGuardrail) -> bool: + """Whether a streaming pipeline step can run this guardrail against the buffered + stream: through the unified path, or through its own post-call hook on the + assembled response. A guardrail with neither (one that only rewrites the stream + through its iterator hook) has to keep running on its own.""" + return ( + PipelineExecutor.supports_unified_execution(callback) + or type(callback).async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + ) + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..bbec1b842c4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -451,7 +451,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) -def _pipeline_managed_guardrail_names( +def pipeline_managed_guardrail_names( data: Mapping[str, object], mode: Literal["pre_call", "post_call"] ) -> frozenset[str]: return _pipeline_step_guardrail_names( @@ -514,9 +514,9 @@ def _merge_pipeline_metadata_writes( _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - return callback is not None and PipelineExecutor.supports_unified_execution(callback) + return callback is not None and PipelineExecutor.supports_streaming_execution(callback) def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -541,14 +541,15 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( - step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail) + step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) ) ) if not unsupported: return True verbose_proxy_logger.warning( - "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " - "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s", + "Policy '%s' has post_call pipeline guardrails with neither the unified apply_guardrail interface nor a " + "post-call hook, one of which streaming pipelines need; the stream skips the pipeline and its guardrails " + "run on their own: %s", policy_name, ", ".join(unsupported), ) @@ -562,11 +563,12 @@ def _streamable_post_call_pipelines( The post_call pipelines a streaming response can be gated through. Streaming pipelines scan the buffered stream through the endpoint guardrail - translation of the request route, so every step's guardrail needs the - unified apply_guardrail interface and the route needs a translation. A - pipeline that cannot be run that way yet is left out and its guardrails - run on the stream on their own, the way they did before pipelines ran on - streams at all, with a warning naming the pipeline. + translation of the request route, so every step's guardrail needs either the + unified apply_guardrail interface or a post-call hook to run against the + assembled response, and the route needs a translation. A pipeline that + cannot be run that way yet is left out and its guardrails run on the stream + on their own, the way they did before pipelines ran on streams at all, with + a warning naming the pipeline. """ post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: @@ -1968,7 +1970,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call") caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in @@ -2956,7 +2958,7 @@ class ProxyLogging: if pipeline_response is not None: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below - pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run @@ -3272,7 +3274,7 @@ class ProxyLogging: _cached_guardrail_data: dict | None = None _guardrail_data_computed = False pipeline_managed: Final = ( - _pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() ) for callback in litellm.callbacks: 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 bf40f781fa3..a10a2aae0ee 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 @@ -2156,3 +2156,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert open_key == StreamingScanKey(texts=("hi",)) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + + +class TestAnthropicMessagesHandlerPostCallHookResponse: + def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + assembled = ModelResponse( + id="msg_1", + model="claude", + choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + ) + + hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled) + + assert hook_response["type"] == "message" + assert hook_response["role"] == "assistant" + assert hook_response["content"] == [{"type": "text", "text": "hello world"}] + assert hook_response["stop_reason"] == "end_turn" + assert hook_response["usage"]["input_tokens"] == 1 + assert hook_response["usage"]["output_tokens"] == 2 + + def test_anything_else_reaches_the_hook_untouched(self): + native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} + + assert AnthropicMessagesHandler().post_call_hook_response(native) is native diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 54ef1f79f4d..db403a17c28 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -534,7 +534,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - monkeypatch.setattr(litellm, "callbacks", []) result = await PipelineExecutor.execute_steps( steps=pipeline.steps, @@ -1153,3 +1152,152 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] + + +class _LegacyHookGuardrail(CustomGuardrail): + """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" + + def __init__(self, replacement=None, raises=None): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.replacement = replacement + self.raises = raises + self.calls = [] + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) + if self.raises is not None: + raise self.raises + return self.replacement + + +class _NativeHooksGuardrail(_LegacyHookGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + +class _LegacyScanningTranslation: + """Stores the assembled response under request_data["response"] before scanning, like the + chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one + text per entry of a replacement's "texts".""" + + delivers_ended_stream_text_rewrites = True + + def post_call_hook_response(self, response): + return {"native": True, "text": response["text"]} + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) + outputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": list(response["texts"])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + return await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail=on_fail, on_error=on_error)], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=chunks, + endpoint_translation=_LegacyScanningTranslation(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) +async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): + guardrail = guardrail_class(replacement={"texts": ["[REWRITTEN] hello world"]}) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in guardrail.calls] == [{"native": True, "text": "hello world"}] + assert guardrail.calls[0]["data"]["model"] == "m" + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert len(guardrail.calls) == 1 + assert chunks == [_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch): + exc = HTTPException(status_code=400, detail={"error": "output blocked"}) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["fail"] + assert result.original_exception is exc + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch): + chunks = [_chunk()] + + result = await _run_legacy_streaming_step( + monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block" + ) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["error"] + assert result.step_results[0].error_detail == "boom" + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement={"texts": ["split", "in two"]}) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..f4165f98ba5 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -671,6 +671,32 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp assert routed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch): + """A post_call pipeline step already ran the opted-out guardrail's own hook against + the buffered stream, so the deferred audit must not run it a second time.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == [] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 73dd6746b29..d01f881c446 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json +from copy import deepcopy import logging from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -1497,29 +1498,78 @@ async def _async_chunk_iter(chunks: List[Any]): yield chunk -def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( +def _legacy_hook_stream_guardrail( + seen: Dict[str, Any], + rewrite: Callable[[Any], Any] | None = None, + raises: Exception | None = None, + native_lifecycle: bool = False, +) -> CustomGuardrail: + class LegacyHookGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = native_lifecycle + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["data"] = data + seen["user_api_key_dict"] = user_api_key_dict + seen["response"] = deepcopy(response) + if raises is not None: + raise raises + return None if rewrite is None else rewrite(response) + + if native_lifecycle: + + class NativeLifecycleGuardrail(LegacyHookGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + return NativeLifecycleGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + +def _rewritten_model_response(response: Any) -> litellm.ModelResponse: + payload = response.model_dump() + payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"] + return litellm.ModelResponse(**payload) + + +def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only( make_user_api_key_auth, monkeypatch, caplog ): - class NativeOnlyGuardrail(CustomGuardrail): - pass - supported = _unified_stream_guardrail({}) - native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call) - monkeypatch.setattr(litellm, "callbacks", [supported, native_only]) - governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + legacy = _legacy_hook_stream_guardrail({}) + legacy.guardrail_name = "gr-legacy" + iterator_only = _iterator_hook_only_guardrail("gr-iterator", {}) + monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only]) + governed = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")], + ) ungoverned = GuardrailPipeline( mode="post_call", - steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")], ) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")]) data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) assert streamable == (("governed", governed),) - assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog)) - assert not any("'governed'" in message for message in _warnings(caplog)) + assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog)) + assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( @@ -1569,55 +1619,123 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u @pytest.mark.asyncio @pytest.mark.parametrize("native_lifecycle", [False, True]) -async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support( +async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog ): seen: Dict[str, Any] = {} - if native_lifecycle: - - class NativeOnlyGuardrail(CustomGuardrail): - use_native_lifecycle_hooks = True - - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): - seen["count"] = seen.get("count", 0) + 1 - return inputs - - else: - - class NativeOnlyGuardrail(CustomGuardrail): - async def async_post_call_success_hook(self, data, user_api_key_dict, response): - seen["count"] = seen.get("count", 0) + 1 - return response - - monkeypatch.setattr( - litellm, - "callbacks", - [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], - ) + guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) chunks = _stream_chunks() - delivered: List[Any] = [] + auth = make_user_api_key_auth(request_route="/v1/chat/completions") with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="completion", - guardrails_only=True, + user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert out is not None and out.get("stream") is True + assert seen["count"] == 1 + assert isinstance(seen["response"], litellm.ModelResponse) + assert seen["response"].choices[0].message.content == "hello world" + assert seen["data"]["messages"] == data["messages"] + assert seen["user_api_key_dict"] is auth + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), response=_async_chunk_iter(chunks), request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + blocked = HTTPException(status_code=400, detail={"error": "output blocked"}) + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, ): delivered.append(item) - assert out is not None - assert out.get("stream") is True - assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] - assert len(delivered) == 2 - assert seen.get("count") is None - assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + with pytest.raises(HTTPException) as info: + await _drain() + + assert seen["count"] == 1 + assert delivered == [] + assert info.value is blocked + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + def rewrite(response: Any) -> Dict[str, Any]: + return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]} + + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_sse_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert seen["response"]["content"][0]["text"] == "hello world" + assert seen["response"]["role"] == "assistant" + raw = b"".join(delivered).decode() + assert "[REWRITTEN] hello world" in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw @pytest.mark.asyncio @@ -1625,19 +1743,7 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} - - class IteratorHookGuardrail(CustomGuardrail): - async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): - seen["count"] = seen.get("count", 0) + 1 - async for item in response: - item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" - yield item - - monkeypatch.setattr( - litellm, - "callbacks", - [IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], - ) + monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) From d497434680186367ce42795ae1fca314503cd2aa Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 8 Sep 2026 19:28:04 +0000 Subject: [PATCH 039/157] fix(model_prices): update provider and web search pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 140 +++++++++++++++--- model_prices_and_context_window.json | 140 +++++++++++++++--- .../test_tool_call_cost_tracking.py | 65 +++++++- ...tex_ai_gemini_transcribe_transformation.py | 4 +- 4 files changed, 293 insertions(+), 56 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ef8da798922..a540df1c7d1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29176,7 +29176,13 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_web_search": true }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29193,9 +29199,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29204,7 +29210,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4o-mini-audio-preview": { "deprecation_date": "2026-05-07", @@ -29294,9 +29301,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29321,9 +29328,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29434,9 +29441,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29461,9 +29468,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -41625,6 +41632,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -47811,7 +47840,7 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, + "max_input_tokens": 128000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", @@ -47827,7 +47856,7 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, + "max_input_tokens": 128000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", @@ -47842,14 +47871,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47858,14 +47890,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47873,6 +47908,42 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -56824,8 +56895,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56860,6 +56931,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ef8da798922..a540df1c7d1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29176,7 +29176,13 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, + "supports_web_search": true }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29193,9 +29199,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29204,7 +29210,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4o-mini-audio-preview": { "deprecation_date": "2026-05-07", @@ -29294,9 +29301,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29321,9 +29328,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29434,9 +29441,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29461,9 +29468,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -41625,6 +41632,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -47811,7 +47840,7 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, + "max_input_tokens": 128000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", @@ -47827,7 +47856,7 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, + "max_input_tokens": 128000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", @@ -47842,14 +47871,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47858,14 +47890,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47873,6 +47908,42 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -56824,8 +56895,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56860,6 +56931,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 6cc3dcceebc..38bf8c544c7 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,5 +1,6 @@ -import os +import json from collections.abc import Mapping, Sequence +from pathlib import Path import pytest @@ -11,8 +12,6 @@ from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebS from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams - - def test_web_search_cost_low(): web_search_options = WebSearchOptions(search_context_size="low") model_info = litellm.get_model_info("gpt-4o-search-preview") @@ -683,12 +682,13 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( def _openai_responses_with_web_search_calls(model, num_calls): - from litellm.types.llms.openai import ResponsesAPIResponse from openai.types.responses.response_function_web_search import ( ActionSearch, ResponseFunctionWebSearch, ) + from litellm.types.llms.openai import ResponsesAPIResponse + output = [ ResponseFunctionWebSearch( id=f"ws_{i}", @@ -859,11 +859,64 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map) custom_llm_provider="openai", standard_built_in_tools_params=None, ) - assert cost == pytest.approx(0.035), ( - f"dated search-preview id must bill the $0.035 search fee, got ${cost}" + assert cost == pytest.approx(0.025), ( + f"dated search-preview id must bill the $0.025 search fee, got ${cost}" ) +@pytest.mark.parametrize( + "web_search_options", + [ + None, + WebSearchOptions(search_context_size="low"), + WebSearchOptions(search_context_size="medium"), + WebSearchOptions(search_context_size="high"), + ], +) +def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( + web_search_options, local_model_cost_map +): + """Snapshot and alias must bill web search identically: OpenAI lists preview search at $25 per 1k calls.""" + alias_info = litellm.get_model_info("gpt-4o-mini") + snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18") + + assert snapshot_info["supports_web_search"] is True + assert alias_info["supports_web_search"] is True + + snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=snapshot_info + ) + alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=alias_info + ) + + assert snapshot_cost == alias_cost == 0.025 + + +def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps(): + """The bundled backup and canonical file are both served to deployments, so both must carry the price.""" + repo_root = Path(__file__).parents[4] + cost_maps = tuple( + json.loads((repo_root / path).read_text(encoding="utf-8")) + for path in ( + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ) + ) + canonical, backup = cost_maps + expected_search_price = { + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025, + "search_context_size_high": 0.025, + } + for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"): + canonical_entry = canonical[model_name] + backup_entry = backup[model_name] + assert canonical_entry["search_context_cost_per_query"] == expected_search_price + assert backup_entry["search_context_cost_per_query"] == expected_search_price + assert canonical_entry == backup_entry + + # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index eadd87d9c92..08e46b1ffac 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -322,8 +322,8 @@ class TestModelCostEntry: entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"] assert entry["mode"] == "audio_transcription" assert entry["litellm_provider"] == "vertex_ai" - assert entry["input_cost_per_audio_token"] == pytest.approx(2.5e-06) - assert entry["input_cost_per_token"] == pytest.approx(2.5e-06) + assert entry["input_cost_per_audio_token"] == pytest.approx(2e-06) + assert entry["input_cost_per_token"] == pytest.approx(2e-06) assert entry["output_cost_per_token"] == pytest.approx(1.2e-05) assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] From 2823d09e1bcfc99cee126583f3552e5d5b32dac5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:30:05 -0700 Subject: [PATCH 040/157] refactor(policy_engine): type the request-state parameter of the retrieval hook --- litellm/proxy/policy_engine/response_retrieval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 54b9acf6150..0aa9579ca71 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -71,7 +71,7 @@ def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[Polic def attach_post_call_pipelines_to_retrieval( - data: dict, # mutable-ok: the proxy's request-state dict, written in place like every other policy engine hook + data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place user_api_key_dict: "UserAPIKeyAuth", llm_router: "Router | None", ) -> None: From 3f30c05493091f62d1edcfb0af58aa33a370a335 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:39:39 -0700 Subject: [PATCH 041/157] fix(guardrails): skip only stream-gated pipeline guardrails in the deferred post-call pass --- litellm/proxy/common_request_processing.py | 6 +- litellm/proxy/utils.py | 37 +++++++++-- .../test_proxy_logging_hook_detection.py | 63 +++++++++++++++++++ 3 files changed, 98 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7e85b66682c..2a6eb0a16ff 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3343,15 +3343,15 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.proxy_server import llm_router as _global_llm_router from litellm.proxy.utils import ( _check_and_merge_model_level_guardrails, - pipeline_managed_guardrail_names, + stream_gated_guardrail_names, ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) - pipeline_managed: Final = pipeline_managed_guardrail_names(captured_data, "post_call") + stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue - if cb.guardrail_name in pipeline_managed: + if cb.guardrail_name in stream_gated: continue if not cb.should_run_guardrail( data=guardrail_data, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bbec1b842c4..106fab7af8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -538,12 +538,16 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = tuple( +def _pipeline_unsupported_streaming_guardrails(pipeline: "GuardrailPipeline") -> tuple[str, ...]: + return tuple( dict.fromkeys( step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) ) ) + + +def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline) if not unsupported: return True verbose_proxy_logger.warning( @@ -573,13 +577,12 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - route: Final = user_api_key_dict.request_route - if route and resolve_endpoint_translation(user_api_key_dict, None) is None: + if not _route_has_endpoint_translation(user_api_key_dict): verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " "on their own: %s", - route, + user_api_key_dict.request_route, ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), ) return () @@ -590,6 +593,30 @@ def _streamable_post_call_pipelines( ) +def _route_has_endpoint_translation(user_api_key_dict: UserAPIKeyAuth) -> bool: + return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None + + +def stream_gated_guardrail_names( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + """ + The guardrails whose post_call pipelines gate a streaming response on this + route: the selection ``_streamable_post_call_pipelines`` makes, without its + warnings, so the post-call pass deferred to the end of the stream skips + exactly the guardrails the pipelines already ran and no others. + """ + if not _route_has_endpoint_translation(user_api_key_dict): + return frozenset() + return _pipeline_step_guardrail_names( + tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline) + ) + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index f4165f98ba5..10de7e7fc3e 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -697,6 +697,69 @@ async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monk assert pipeline_managed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch): + """A pipeline step with neither streaming interface keeps the whole pipeline off the + stream, so the deferred audit is the only place the opted-out guardrail's own hook + still runs, the way it did before pipelines ran on streams.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + class NeitherHookGuardrail(CustomGuardrail): + pass + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither]) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="keeps_native", on_fail="next"), + PipelineStep(guardrail="gr-neither", on_fail="block"), + ], + ) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch): + """A route with no endpoint guardrail translation cannot gate the stream through its + pipelines, so the deferred audit still owes the opted-out guardrail its own hook.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be From 8eb6f6d8350dd0ef7f2e5a61b042213e0b6bd319 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 8 Sep 2026 19:43:32 +0000 Subject: [PATCH 042/157] test(cost_calc): type web search test params and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/test_tool_call_cost_tracking.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 38bf8c544c7..f319413045f 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -874,9 +874,8 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map) ], ) def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( - web_search_options, local_model_cost_map -): - """Snapshot and alias must bill web search identically: OpenAI lists preview search at $25 per 1k calls.""" + web_search_options: WebSearchOptions | None, local_model_cost_map: None +) -> None: alias_info = litellm.get_model_info("gpt-4o-mini") snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18") @@ -894,7 +893,6 @@ def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps(): - """The bundled backup and canonical file are both served to deployments, so both must carry the price.""" repo_root = Path(__file__).parents[4] cost_maps = tuple( json.loads((repo_root / path).read_text(encoding="utf-8")) From 91091fd93e7b7be6576be26e1ef1216391af667e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:59:58 -0700 Subject: [PATCH 043/157] fix: withdraw policy header claims while a background response is pending and log ungoverned retrievals --- .../proxy/policy_engine/response_retrieval.py | 31 ++++-- litellm/proxy/utils.py | 62 +++++++++++- .../policy_engine/test_response_retrieval.py | 36 +++++-- .../proxy_logging/test_guardrail_pipeline.py | 97 +++++++++++++++++++ 4 files changed, 205 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 0aa9579ca71..8fc75e40070 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -1,6 +1,7 @@ from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter @@ -29,14 +30,23 @@ PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] _POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) -def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | None: - if llm_router is None or not isinstance(response_id, str): - return None - model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) +@dataclass(frozen=True, slots=True) +class UngovernedRetrieval: + reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"] + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval: + if llm_router is None: + return UngovernedRetrieval("no router") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None + ) if model_id is None: - return None + return UngovernedRetrieval("response id names no deployment") deployment: Final = llm_router.get_deployment(model_id) - return deployment.model_name if deployment is not None else None + if deployment is None: + return UngovernedRetrieval("deployment no longer in the router") + return deployment.model_name def _retrieval_context( @@ -78,7 +88,12 @@ def attach_post_call_pipelines_to_retrieval( if not get_policy_registry().is_initialized(): return model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) - if model_group is None: + if isinstance(model_group, UngovernedRetrieval): + verbose_proxy_logger.warning( + "Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)", + data.get("response_id"), + model_group.reason, + ) return context: Final = _retrieval_context(data, user_api_key_dict, model_group) post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3d03d3858ee..1972d92b610 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -93,6 +93,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, + get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, ) @@ -156,6 +157,8 @@ from litellm.proxy.hooks.sensitive_data_routing import ( from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -533,16 +536,65 @@ def _is_pending_background_response(response: LLMResponseTypes) -> bool: return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES -def _log_deferred_post_call_pipelines(data: Mapping[str, object], response: ResponsesAPIResponse) -> None: - policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) - if not policy_names: +def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]: + resolved: Final = PolicyResolver.resolve_policy_guardrails( + policy_name=policy_name, policies=get_policy_registry().get_all_policies() + ) + return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) + + +def _without_names( + bucket: dict, # mutable-ok: the applied_* header slots live in the request-state dict every hook writes in place + slot: str, + names: frozenset[str], +) -> None: + claimed: Final = bucket.get(slot) + if not isinstance(claimed, list): + return + remaining: Final = [name for name in claimed if name not in names] + if remaining: + bucket[slot] = remaining + else: + bucket.pop(slot) + + +def _withdraw_deferred_claims( + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + deferred: Sequence[tuple[str, "GuardrailPipeline"]], +) -> None: + outside_by_policy: Final = MappingProxyType( + {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} + ) + running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(*outside_by_policy.values()) + withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) + withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere + _, bucket = get_or_create_metadata_bucket(data) + _without_names(bucket, "applied_policies", withdrawn_policies) + _without_names(bucket, "applied_guardrails", withdrawn_guardrails) + sources: Final = bucket.get("policy_sources") + if not isinstance(sources, dict): + return + remaining_sources: Final = {name: reason for name, reason in sources.items() if name not in withdrawn_policies} + if remaining_sources: + bucket["policy_sources"] = remaining_sources + else: + bucket.pop("policy_sources") + + +def _defer_post_call_pipelines( + data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + response: ResponsesAPIResponse, +) -> None: + deferred: Final = _post_call_pipelines(data) + if not deferred: return verbose_proxy_logger.debug( "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", response.id, response.status, - ", ".join(policy_names), + ", ".join(policy_name for policy_name, _pipeline in deferred), ) + _withdraw_deferred_claims(data, deferred) def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: @@ -2939,7 +2991,7 @@ class ProxyLogging: response: LLMResponseTypes, ) -> LLMResponseTypes | None: if _is_pending_background_response(response): - _log_deferred_post_call_pipelines(data, response) + _defer_post_call_pipelines(data, response) return None _, pipeline_response = await self._maybe_execute_pipelines( data=data, diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index f7f29d6b46e..8c00b18eb63 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,3 +1,5 @@ +import logging + import pytest from litellm.proxy._types import UserAPIKeyAuth @@ -128,30 +130,48 @@ def test_already_attached_policy_is_not_attached_twice(policy_engine): assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] +def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "retrieved without its post_call policy pipelines" in record.getMessage() + ] + + @pytest.mark.parametrize( - "response_id", - ["resp_plain_upstream_id", _encoded_response_id("deployment-missing-from-router"), None], + ("response_id", "reason"), + [ + ("resp_plain_upstream_id", "response id names no deployment"), + (_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"), + (None, "response id names no deployment"), + ], ) -def test_unresolvable_response_id_attaches_nothing(policy_engine, response_id): +def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, caplog, response_id, reason): data = {"response_id": response_id, "litellm_metadata": {}} - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) assert data == {"response_id": response_id, "litellm_metadata": {}} + assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_a_router_attaches_nothing(policy_engine): +def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): data = _retrieval_data(GOVERNED_MODEL_ID) - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_policy_engine_attaches_nothing(): +def test_without_policy_engine_attaches_nothing_quietly(caplog): get_policy_registry().clear() data = _retrieval_data(GOVERNED_MODEL_ID) - attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert _ungoverned_retrieval_warnings(caplog) == [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 7d2c232ad97..8da07f39fd8 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1481,6 +1481,103 @@ async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_resp assert seen["response"] is response +def _output_passing_callbacks() -> list[CustomGuardrail]: + class OutputPassingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + + +def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} + get_policy_registry().load_policies( + { + policy_name: { + "guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]}, + "pipeline": {"mode": "post_call", "steps": [step]}, + } + for policy_name in policy_names + } + ) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names], + "_pipeline_managed_guardrails": {"gr-post"}, + "applied_policies": list(policy_names), + "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], + "policy_sources": {policy_name: "model:m" for policy_name in policy_names}, + }, + } + + +@pytest.fixture +def clear_policy_registry(): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + yield + get_policy_registry().clear() + + +@pytest.mark.asyncio +async def test_pending_background_response_withdraws_the_deferred_policy_claims( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "applied_policies" not in data["metadata"] + assert "policy_sources" not in data["metadata"] + assert "applied_guardrails" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data( + "input-and-output-governance", "response-governance", extra_guardrails={"input-and-output-governance": ["gr-pre"]} + ) + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["input-and-output-governance"] + assert data["metadata"]["applied_guardrails"] == ["gr-pre"] + assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} + + +@pytest.mark.asyncio +async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["response-governance"] + assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"} + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + @pytest.mark.asyncio async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, caplog From 8ce3fe4fcf5ba43278c4eb7ec09c58df513d36c5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:13:36 -0700 Subject: [PATCH 044/157] style: give the header slot rebuilds and the in-place slot write their lint reasons --- litellm/proxy/utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1972d92b610..e48bd3851fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -551,9 +551,11 @@ def _without_names( claimed: Final = bucket.get(slot) if not isinstance(claimed, list): return - remaining: Final = [name for name in claimed if name not in names] + remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to + name for name in claimed if name not in names + ] if remaining: - bucket[slot] = remaining + bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place else: bucket.pop(slot) @@ -574,7 +576,9 @@ def _withdraw_deferred_claims( sources: Final = bucket.get("policy_sources") if not isinstance(sources, dict): return - remaining_sources: Final = {name: reason for name, reason in sources.items() if name not in withdrawn_policies} + remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place + name: reason for name, reason in sources.items() if name not in withdrawn_policies + } if remaining_sources: bucket["policy_sources"] = remaining_sources else: From f59354d09b642bf50d888a82b5fac5b96d34169b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:22:29 -0700 Subject: [PATCH 045/157] refactor(guardrails): patch Anthropic SSE rewrites by field The ended-stream rewriters now describe the one field they change as a frozen _SSEFieldRewrite and a single applier builds the patched event, so the handler adds no mutable-collection constructions over staging's total --- .../chat/guardrail_translation/handler.py | 42 +++++++++++++------ 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2049866b444..cbbccac17f3 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,7 +13,7 @@ Pattern Overview: """ import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat @@ -161,7 +161,25 @@ class _ToolCallShape: arguments: str -_SSEEventRewriter = Callable[[Mapping[str, object]], Mapping[str, object] | None] +@dataclass(frozen=True, slots=True) +class _SSEFieldRewrite: + """One field of one nested section of a buffered SSE event, rewritten.""" + + section: str + field: str + value: object + + +class _SSEEventRewriter(Protocol): + def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ... + + +def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]: + rewrite: Final = rewrite_event(event) + section: Final = None if rewrite is None else event.get(rewrite.section) + if rewrite is None or not isinstance(section, Mapping): + return event + return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: @@ -1253,13 +1271,13 @@ class AnthropicMessagesHandler(BaseTranslation): message and content-block framing untouched.""" replacements: Final = chain((rewritten_text,), repeat("")) - def rewrite_text_delta(event: Mapping[str, object]) -> Mapping[str, object] | None: + def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None: delta: Final = event.get("delta") if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): return None if delta.get("type") != "text_delta": return None - return {**event, "delta": {**delta, "text": next(replacements)}} + return _SSEFieldRewrite("delta", "text", next(replacements)) AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) @@ -1306,22 +1324,21 @@ class AnthropicMessagesHandler(BaseTranslation): {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} ) - def rewrite_tool_use(event: Mapping[str, object]) -> Mapping[str, object] | None: + def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None: index: Final = event.get("index") if not isinstance(index, int) or index not in rewrites_by_block: return None match event.get("type"): case "content_block_start": - block: Final = event.get("content_block") name: Final = rewrites_by_block[index].name - if not isinstance(block, Mapping) or name is None: + if name is None: return None - return {**event, "content_block": {**block, "name": name}} + return _SSEFieldRewrite("content_block", "name", name) case "content_block_delta": delta: Final = event.get("delta") if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": return None - return {**event, "delta": {**delta, "partial_json": next(argument_replacements[index])}} + return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index])) case _: return None @@ -1343,8 +1360,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: if isinstance(item, dict): - rewritten: Final = rewrite_event(_as_str_mapping(item)) - return item if rewritten is None else dict(rewritten) + return _rewritten_event(_as_str_mapping(item), rewrite_event) if isinstance(item, (bytes, bytearray)): return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) return item @@ -1374,8 +1390,8 @@ class AnthropicMessagesHandler(BaseTranslation): return line if not isinstance(data, dict): return line - rewritten: Final = rewrite_event(_as_str_mapping(data)) - return line if rewritten is None else "data: " + json.dumps(rewritten) + rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) From 89e11949c834d0b466375b0442b27304946568a9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:22:30 -0700 Subject: [PATCH 046/157] fix(guardrails): key Responses stream tool-call rewrites by call_id Bridged Responses streams give reasoning and message items output_index 0 and start function calls at 1, so keying rewrites by output_index rewrote the wrong items. Rewrites now follow each function call's call_id through the buffered item and argument events, refuse when an event cannot be resolved to a rewritten call, and the refusal branches on all three handlers get regression tests --- .../guardrail_translation/handler.py | 132 ++++++++++++------ .../test_anthropic_guardrail_handler.py | 25 ++++ .../test_openai_guardrail_handler.py | 56 ++++++++ ...test_openai_responses_guardrail_handler.py | 90 ++++++++++++ 4 files changed, 259 insertions(+), 44 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 208ceb959e6..280a8670c36 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -140,6 +140,10 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( ) +_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset( + {"response.function_call_arguments.delta", "response.function_call_arguments.done"} +) +_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call_output": "output", "message": "content"} ) @@ -930,70 +934,110 @@ class OpenAIResponsesHandler(BaseTranslation): ) -> None: """Write ended-stream guardrail tool-call rewrites into the completed envelope's ``function_call`` items and sync the earlier stream events, - keyed by ``output_index``. The guardrail sees the envelope's function - calls in output order, which is how a rewritten call finds its item; a - rewrite whose calls do not line up with the envelope is reported as - undeliverable, so the pipeline executor discards it and releases the - original events.""" + keyed by ``call_id``. The guardrail sees the envelope's function calls + in output order, which is how a rewritten call finds its ``call_id``; + the stream events find their call through the ``call_id`` on + ``output_item`` events and the ``item_id`` on argument events, since an + event's ``output_index`` need not match the envelope's (the chat bridge + numbers tool calls from 1 while the envelope lists them after the + message). A rewrite whose calls do not line up with the envelope, or + whose events cannot be found, is reported as undeliverable, so the + pipeline executor discards it and releases the original events.""" if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - function_call_indices: Final = tuple( - output_idx - for output_idx, output_item in enumerate(outputs) - if stream_item_field(output_item, "type") == "function_call" + function_call_items: Final = tuple( + output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call" ) - if len(function_call_indices) != len(post_guardrail_tool_calls): - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - - raise UndeliverableStreamRewrite(guardrail_name) - rewrites_by_output_index: Final = MappingProxyType( + call_ids: Final = tuple( + call_id + for output_item in function_call_items + if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id + ) + stream_events: Final = responses_so_far[:-1] + call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events) + event_call_ids: Final = tuple( + self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events + ) + rewrites_by_call_id: Final = MappingProxyType( { - output_idx: after - for output_idx, before, after in zip( - function_call_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls - ) + call_id: after + for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before } ) - for output_idx, rewrite in rewrites_by_output_index.items(): - self._write_function_call_item(outputs[output_idx], rewrite.name, rewrite.arguments) - self._sync_stream_events_with_tool_call_rewrites( - stream_events=responses_so_far[:-1], - rewrites_by_output_index=rewrites_by_output_index, + unresolved_argument_event: Final = any( + call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES + for event, call_id in zip(stream_events, event_call_ids) ) + if ( + len(call_ids) != len(function_call_items) + or len(frozenset(call_ids)) != len(call_ids) + or len(call_ids) != len(post_guardrail_tool_calls) + or unresolved_argument_event + or not rewrites_by_call_id.keys() <= frozenset(event_call_ids) + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - def _sync_stream_events_with_tool_call_rewrites( - self, - stream_events: Sequence[object], - rewrites_by_output_index: Mapping[int, _ToolCallShape], - ) -> None: - """Sync pre-completion function-call events with the rewritten completed - response: the first ``function_call_arguments.delta`` for a rewritten call - carries the full rewritten arguments and the rest are blanked, while - ``function_call_arguments.done`` and ``output_item.done`` carry the full - rewritten arguments and ``output_item.added`` / ``output_item.done`` the - rewritten name, so every event a client may read agrees with the - rewritten ``response.completed`` payload.""" + raise UndeliverableStreamRewrite(guardrail_name) + for output_item, rewrite in ( + (output_item, rewrites_by_call_id[call_id]) + for output_item, call_id in zip(function_call_items, call_ids) + if call_id in rewrites_by_call_id + ): + self._write_function_call_item(output_item, rewrite.name, rewrite.arguments) delta_replacements: Final = MappingProxyType( - {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_output_index.items()} + {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()} ) - for event in stream_events: - output_index = stream_item_field(event, "output_index") - if not isinstance(output_index, int) or output_index not in rewrites_by_output_index: + for event, call_id in zip(stream_events, event_call_ids): + if call_id not in rewrites_by_call_id: continue - rewrite = rewrites_by_output_index[output_index] match stream_item_field(event, "type"): case "response.function_call_arguments.delta": - self._write_event_field(event, "delta", next(delta_replacements[output_index])) + self._write_event_field(event, "delta", next(delta_replacements[call_id])) case "response.function_call_arguments.done": - self._write_event_field(event, "arguments", rewrite.arguments) + self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments) case "response.output_item.added": - self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, None) + self._write_function_call_item( + stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None + ) case "response.output_item.done": - self._write_function_call_item(stream_item_field(event, "item"), rewrite.name, rewrite.arguments) + self._write_function_call_item( + stream_item_field(event, "item"), + rewrites_by_call_id[call_id].name, + rewrites_by_call_id[call_id].arguments, + ) case _: pass + @staticmethod + def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: + items: Final = tuple( + stream_item_field(event, "item") + for event in stream_events + if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + ) + return MappingProxyType( + { + item_id: call_id + for item in items + if stream_item_field(item, "type") == "function_call" + and isinstance(item_id := stream_item_field(item, "id"), str) + and isinstance(call_id := stream_item_field(item, "call_id"), str) + } + ) + + @staticmethod + def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: + event_type: Final = stream_item_field(event, "type") + if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES: + item_id: Final = stream_item_field(event, "item_id") + return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None + if event_type not in _OUTPUT_ITEM_EVENT_TYPES: + return None + item: Final = stream_item_field(event, "item") + call_id: Final = stream_item_field(item, "call_id") + return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None + @staticmethod def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: if item is None: 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 bd2dff9b33c..64c243343d5 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 @@ -381,6 +381,31 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert chunks == original + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + server_tool_use = [ + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ] + tool_use = self._ended_tool_use_sse_chunks() + chunks = ( + tool_use[:1] + + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use] + + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]] + ) + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index be168ec83e6..c7011a6f00e 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1252,6 +1252,62 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _two_choice_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk( + choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None + ) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=choice_index, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), + chunk(0, fragment('{"fruit": "persimmon"}')), + chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, None, finish_reason="tool_calls"), + chunk(1, None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index bdf7f83757c..fa1e969b277 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1314,6 +1314,96 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' assert completed_event.response.output[0].name == "lookup_fruit" + @staticmethod + def _bridged_function_call_stream_events() -> List[dict]: + reasoning = {"type": "reasoning", "id": "rs_1", "summary": []} + text = {"type": "output_text", "text": "Looking that up", "annotations": []} + message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]} + + def function_call(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"}, + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}}, + {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_1", + "output_index": 1, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "claude-haiku-4-5", + "output": [ + dict(reasoning), + {**message, "content": [dict(text)]}, + function_call('{"fruit": "persimmon"}', "completed"), + ], + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self): + handler = OpenAIResponsesHandler() + events = self._bridged_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[6]["delta"] == '{"fruit": "[MASKED]"}' + assert events[7]["delta"] == "" + assert events[8]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["item"]["name"] == "lookup_fruit" + assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[3]["delta"] == "Looking that up" + assert events[4]["item"]["content"][0]["text"] == "Looking that up" + assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up" + assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []} + + @pytest.mark.asyncio + @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"]) + async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + envelope_item = events[5]["response"]["output"][0] + if mismatch == "orphan_call_id": + events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}] + else: + events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): handler = OpenAIResponsesHandler() From 071f83980c319b9d9a21f77179781052c050b517 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:26:17 -0700 Subject: [PATCH 047/157] refactor(proxy): type the post_call deferral helpers' request-state parameters --- litellm/proxy/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e48bd3851fb..e653ca6b62c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -544,7 +544,7 @@ def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline" def _without_names( - bucket: dict, # mutable-ok: the applied_* header slots live in the request-state dict every hook writes in place + bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write slot: str, names: frozenset[str], ) -> None: @@ -561,7 +561,7 @@ def _without_names( def _withdraw_deferred_claims( - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data deferred: Sequence[tuple[str, "GuardrailPipeline"]], ) -> None: outside_by_policy: Final = MappingProxyType( @@ -586,7 +586,7 @@ def _withdraw_deferred_claims( def _defer_post_call_pipelines( - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data response: ResponsesAPIResponse, ) -> None: deferred: Final = _post_call_pipelines(data) @@ -2990,7 +2990,7 @@ class ProxyLogging: async def _run_post_call_pipelines( self, - data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes, ) -> LLMResponseTypes | None: From c1132d045470853a8a186d730678493f3f015dbe Mon Sep 17 00:00:00 2001 From: shotsan Date: Tue, 8 Sep 2026 13:44:50 -0700 Subject: [PATCH 048/157] fix(convert_dict_to_response): handle empty choices list without raising 500 APIError (Fixes #40276) --- .../convert_dict_to_response.py | 6 +- .../test_convert_dict_to_response.py | 70 ++++++++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index a6f10e1ede3..7d8072ba622 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -179,7 +179,7 @@ async def convert_to_streaming_response_async( choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): from litellm.exceptions import APIError raise APIError( @@ -287,7 +287,7 @@ def convert_to_streaming_response( model_response_object: Final = ModelResponseStream() choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): from litellm.exceptions import APIError raise APIError( @@ -623,7 +623,7 @@ def convert_to_model_response_object( return convert_to_streaming_response(response_object=response_object) choice_list: Final[list[Choices]] = [] - if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable): + if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): from litellm.exceptions import APIError raise APIError( diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 304d732c518..5406df691aa 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,4 +1,4 @@ - +import pytest from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -99,3 +99,71 @@ def test_handle_invalid_parallel_tool_calls_skips_custom_tool_calls(): ) result = _handle_invalid_parallel_tool_calls([custom_tool_call, function_tool_call]) assert result == [custom_tool_call, function_tool_call] + + +def test_convert_empty_choices_response() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + ) + + resp = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + "vertex_ai_safety_results": ["blocked"], + } + result = convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert result.choices == [] + assert getattr(result, "vertex_ai_safety_results") == ["blocked"] + + # Test sync streaming generator handles empty choices + sync_stream = list(convert_to_streaming_response(response_object=resp)) + assert len(sync_stream) == 1 + assert sync_stream[0].choices == [] + + +@pytest.mark.asyncio +async def test_convert_empty_choices_response_async() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async, + ) + + resp = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + } + async_chunks = [] + async for chunk in convert_to_streaming_response_async(response_object=resp): + async_chunks.append(chunk) + assert len(async_chunks) == 1 + assert async_chunks[0].choices == [] + + +def test_convert_missing_choices_raises_api_error() -> None: + from litellm.exceptions import APIError + + resp = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + } + with pytest.raises(APIError) as exc_info: + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert "no 'choices'" in str(exc_info.value) + From d72eb4491ab2da2a56944c720a3200280cdd040d Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 8 Sep 2026 17:22:30 -0400 Subject: [PATCH 049/157] fix(ui): repair pass-through delete confirm dialog and disable delete for config endpoints --- .../PassThroughEndpointsTable.test.tsx | 47 ++++++++++++++ .../PassThroughEndpointsTableColumns.tsx | 30 +++++++-- .../PassThroughSettings.tsx | 62 ++++++++----------- 3 files changed, 98 insertions(+), 41 deletions(-) diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx index 040fa463f9c..7a3fc621fc1 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx @@ -89,6 +89,53 @@ describe("PassThroughEndpointsTable", () => { expect(onDeleteClick).toHaveBeenCalledWith("ep-1"); }); + it("should disable edit and delete for config-defined endpoints", async () => { + const user = userEvent.setup(); + const onEndpointClick = vi.fn(); + const onDeleteClick = vi.fn(); + const configEndpoint: passThroughItem = { + id: "ep-config", + path: "/from-config", + target: "https://config.example.com", + headers: {}, + is_from_config: true, + }; + render( + , + ); + + await user.click(screen.getByTestId("endpoint-actions-ep-config")); + const editItem = await screen.findByTestId("endpoint-action-edit"); + const deleteItem = await screen.findByTestId("endpoint-action-delete"); + + expect(editItem).toHaveAttribute("data-disabled"); + expect(deleteItem).toHaveAttribute("data-disabled"); + + await user.click(editItem); + await user.click(deleteItem); + + expect(onEndpointClick).not.toHaveBeenCalled(); + expect(onDeleteClick).not.toHaveBeenCalled(); + }); + + it("should label endpoint source as Config or DB", () => { + const configEndpoint: passThroughItem = { + id: "ep-config", + path: "/from-config", + target: "https://config.example.com", + headers: {}, + is_from_config: true, + }; + render(); + expect(screen.getByText("Config")).toBeInTheDocument(); + expect(screen.getAllByText("DB")).toHaveLength(2); + }); + it("should disable edit and delete for endpoints without an id", async () => { const user = userEvent.setup(); const onEndpointClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx index d22b274861a..d63be628110 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx @@ -18,6 +18,10 @@ import { cn } from "@/lib/cva.config"; import type { passThroughItem } from "./PassThroughSettings"; +const CONFIG_EDIT_HINT = "Config pass-through endpoints cannot be edited on the dashboard. Please edit the config file."; +const CONFIG_DELETE_HINT = + "Config pass-through endpoints cannot be deleted on the dashboard. Please edit the config file."; + function HeaderWithTooltip({ title, tooltip }: { title: string; tooltip: string }) { return (
@@ -73,6 +77,7 @@ interface EndpointRowActionsProps { function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: EndpointRowActionsProps) { const endpointId = endpoint.id; + const isFromConfig = endpoint.is_from_config ?? false; return ( endpointId && onEndpointClick(endpointId)} + disabled={isFromConfig || !endpointId} + title={isFromConfig ? CONFIG_EDIT_HINT : undefined} + onClick={() => !isFromConfig && endpointId && onEndpointClick(endpointId)} > Edit @@ -95,8 +101,9 @@ function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: Endpoi endpointId && onDeleteClick(endpointId)} + disabled={isFromConfig || !endpointId} + title={isFromConfig ? CONFIG_DELETE_HINT : undefined} + onClick={() => !isFromConfig && endpointId && onDeleteClick(endpointId)} > Delete @@ -124,7 +131,9 @@ export const getPassThroughEndpointsTableColumns = ({ enableSorting: false, cell: ({ row }) => { const endpointId = row.original.id; - if (!endpointId) return —; + if (!endpointId || row.original.is_from_config) { + return —; + } return ( { + const isFromConfig = row.original.is_from_config ?? false; + return ; + }, + }, { id: "path", accessorKey: "path", diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx index 2dc6fdbd32c..8ef2766d412 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx @@ -1,4 +1,13 @@ import React, { useState, useEffect } from "react"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { deletePassThroughEndpointsCall, getPassThroughEndpointsCall } from "../networking"; import AddPassThroughEndpoint from "../add_pass_through"; @@ -25,6 +34,7 @@ export interface passThroughItem { methods?: string[]; guardrails?: Record; default_query_params?: Record; + is_from_config?: boolean; } const PassThroughSettings: React.FC = ({ accessToken, userRole, userID, premiumUser }) => { @@ -133,42 +143,22 @@ const PassThroughSettings: React.FC = ({ accessToken, onDeleteClick={handleDelete} /> - {isDeleteModalOpen && ( -
-
- - - - -
-
-
-
-

Delete Pass-Through Endpoint

-
-

- Are you sure you want to delete this pass-through endpoint? This action cannot be undone. -

-
-
-
-
-
- - -
-
-
-
- )} + !open && cancelDelete()}> + + + Delete Pass-Through Endpoint + + Are you sure you want to delete this pass-through endpoint? This action cannot be undone. + + + + Cancel + + + +
); }; From 093b473710313f22efd864e5262e5a376a0ee2ce Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 8 Sep 2026 17:43:23 -0400 Subject: [PATCH 050/157] fix(ui): surface the config-endpoint hint as visible menu text --- .../PassThroughEndpointsTable.test.tsx | 12 ++++++++++++ .../PassThroughEndpointsTableColumns.tsx | 12 +++++++----- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx index 7a3fc621fc1..af3f98b746b 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx @@ -115,6 +115,9 @@ describe("PassThroughEndpointsTable", () => { expect(editItem).toHaveAttribute("data-disabled"); expect(deleteItem).toHaveAttribute("data-disabled"); + expect(screen.getByTestId("endpoint-config-hint")).toHaveTextContent( + "This endpoint is defined in the config file and cannot be edited or deleted on the dashboard.", + ); await user.click(editItem); await user.click(deleteItem); @@ -123,6 +126,15 @@ describe("PassThroughEndpointsTable", () => { expect(onDeleteClick).not.toHaveBeenCalled(); }); + it("should not show the config hint for DB endpoints", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByTestId("endpoint-actions-ep-1")); + await screen.findByTestId("endpoint-action-delete"); + expect(screen.queryByTestId("endpoint-config-hint")).not.toBeInTheDocument(); + }); + it("should label endpoint source as Config or DB", () => { const configEndpoint: passThroughItem = { id: "ep-config", diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx index d63be628110..5b18685a140 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx @@ -18,9 +18,8 @@ import { cn } from "@/lib/cva.config"; import type { passThroughItem } from "./PassThroughSettings"; -const CONFIG_EDIT_HINT = "Config pass-through endpoints cannot be edited on the dashboard. Please edit the config file."; -const CONFIG_DELETE_HINT = - "Config pass-through endpoints cannot be deleted on the dashboard. Please edit the config file."; +const CONFIG_ENDPOINT_HINT = + "This endpoint is defined in the config file and cannot be edited or deleted on the dashboard."; function HeaderWithTooltip({ title, tooltip }: { title: string; tooltip: string }) { return ( @@ -91,7 +90,6 @@ function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: Endpoi !isFromConfig && endpointId && onEndpointClick(endpointId)} > @@ -102,12 +100,16 @@ function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: Endpoi variant="destructive" data-testid="endpoint-action-delete" disabled={isFromConfig || !endpointId} - title={isFromConfig ? CONFIG_DELETE_HINT : undefined} onClick={() => !isFromConfig && endpointId && onDeleteClick(endpointId)} > Delete + {isFromConfig && ( +
+ {CONFIG_ENDPOINT_HINT} +
+ )} ); From 547f81c1a52722f487597d93975df1fa79e9464e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:46:23 -0700 Subject: [PATCH 051/157] test: type the background response retrieval test helpers --- .../proxy/policy_engine/test_response_retrieval.py | 9 ++++++--- .../test_litellm/proxy/test_common_request_processing.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 8c00b18eb63..b704f666646 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Mapping import pytest @@ -76,14 +77,16 @@ def policy_engine(): attachment_registry.clear() -def _retrieval_data(model_id: str) -> dict: +def _retrieval_data(model_id: str) -> dict[str, object]: return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} -def _attached_pipelines(data: dict) -> tuple[tuple[str, str], ...]: +def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + bucket = data["litellm_metadata"] + assert isinstance(bucket, dict) return tuple( (policy_name, ",".join(step.guardrail for step in pipeline.steps)) - for policy_name, pipeline in data["litellm_metadata"]["_guardrail_pipelines"] + for policy_name, pipeline in bucket["_guardrail_pipelines"] ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 08007e22cfb..c451a3b4cb0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8340,7 +8340,7 @@ class TestBackgroundResponseRetrievalGovernance: ) return router - async def _pre_call(self, route_type: str, monkeypatch) -> dict: + async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: from litellm.responses.utils import ResponsesAPIRequestUtils client_facing_response_id = "resp_opaque-client-facing-id" From e700e79cce9dac5d0fda0b5d31210283230c9de7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:52:38 -0700 Subject: [PATCH 052/157] refactor(guardrails): keep the rescanned inputs instead of rebuilding the texts list --- litellm/proxy/policy_engine/pipeline_executor.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index afc9a29afd9..089f0683823 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -128,7 +128,7 @@ class _StreamRewriteObserver(CustomGuardrail): class _ScannedTextRecorder(CustomGuardrail): def __init__(self, guardrail_name: str) -> None: super().__init__(guardrail_name=guardrail_name) - self.texts: tuple[str, ...] | None = None + self.inputs: GenericGuardrailAPIInputs | None = None @_logged_by_inner_guardrail async def apply_guardrail( @@ -138,7 +138,7 @@ class _ScannedTextRecorder(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: - self.texts = _text_snapshot(inputs.get("texts")) + self.inputs = inputs return inputs @@ -183,12 +183,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail): if replacement is None: return inputs scanned: Final = _text_snapshot(inputs.get("texts")) - rewritten: Final = await self._scanned_texts(replacement, logging_obj) + rescanned: Final = await self._rescan(replacement, logging_obj) + rewritten: Final = None if rescanned is None else rescanned.get("texts") if scanned is None or rewritten is None or len(rewritten) != len(scanned): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") - return {**inputs, "texts": list(rewritten)} + rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} + return rewritten_inputs - async def _scanned_texts(self, response: object, logging_obj: "LiteLLMLoggingObj | None") -> tuple[str, ...] | None: + async def _rescan( + self, response: object, logging_obj: "LiteLLMLoggingObj | None" + ) -> GenericGuardrailAPIInputs | None: recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") await self.endpoint_translation.process_output_response( response=response, @@ -196,7 +200,7 @@ class _LegacyHookStreamAdapter(CustomGuardrail): litellm_logging_obj=logging_obj, user_api_key_dict=self.user_api_key_dict, ) - return recorder.texts + return recorder.inputs def _prepare_hook_input( From 345298f3c90871a7099952cb08e1e19d396883ee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:04:19 -0700 Subject: [PATCH 053/157] fix(policy_engine): warn when a poll cannot re-match the submitted model name and keep default_on pre_call claims --- .../proxy/policy_engine/response_retrieval.py | 21 +++++++ litellm/proxy/utils.py | 14 ++++- .../policy_engine/test_response_retrieval.py | 56 ++++++++++++++++++- .../proxy_logging/test_guardrail_pipeline.py | 24 ++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 8fc75e40070..d284c44397e 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -18,6 +18,7 @@ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.proxy.policy_engine import PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline @@ -46,9 +47,29 @@ def _model_group_for_response_id(response_id: object, llm_router: "Router | None deployment: Final = llm_router.get_deployment(model_id) if deployment is None: return UngovernedRetrieval("deployment no longer in the router") + hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) + if hidden_by is not None: + verbose_proxy_logger.warning( + "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " + "so a policy attached to the model name it was submitted as does not run on it", + response_id, + deployment.model_name, + hidden_by, + ) return deployment.model_name +def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: + if "*" in model_group: + return "a wildcard deployment" + aliases: Final = tuple( + alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group + ) + if not aliases: + return None + return f"the target of model_group_alias {', '.join(aliases)}" + + def _retrieval_context( data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str ) -> PolicyMatchContext: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e653ca6b62c..37ca52d6fc9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -543,6 +543,16 @@ def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline" return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) +def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]: + return frozenset( + callback.guardrail_name + for callback in litellm.callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + ) + + def _without_names( bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write slot: str, @@ -567,7 +577,9 @@ def _withdraw_deferred_claims( outside_by_policy: Final = MappingProxyType( {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} ) - running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(*outside_by_policy.values()) + running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union( + _guardrails_run_standalone_pre_call(data), *outside_by_policy.values() + ) withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere _, bucket = get_or_create_metadata_bucket(data) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index b704f666646..9dcbef888ce 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -14,11 +14,14 @@ GOVERNED_MODEL_GROUP = "gpt-5.4-mini" GOVERNED_MODEL_ID = "deployment-governed" UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" UNGOVERNED_MODEL_ID = "deployment-ungoverned" +WILDCARD_MODEL_GROUP = "openai/*" +WILDCARD_MODEL_ID = "deployment-wildcard" class FakeRouter: - def __init__(self, deployments: dict[str, Deployment]): + def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None): self._deployments = deployments + self.model_group_alias = model_group_alias or {} def get_deployment(self, model_id: str) -> Deployment | None: return self._deployments.get(model_id) @@ -32,12 +35,14 @@ def _deployment(model_group: str, model_id: str) -> Deployment: ) -def _router() -> FakeRouter: +def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter: return FakeRouter( { GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), - } + WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID), + }, + model_group_alias, ) @@ -133,6 +138,51 @@ def test_already_attached_policy_is_not_attached_twice(policy_engine): assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] +def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage() + ] + + +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(policy_engine, caplog): + data = _retrieval_data(WILDCARD_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(WILDCARD_MODEL_ID) + assert [ + "as model group openai/* (a wildcard deployment)" in message for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_engine, caplog): + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert [ + "(the target of model_group_alias gpt-mini, gpt-hidden)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(policy_engine, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval( + data=_retrieval_data(GOVERNED_MODEL_ID), + user_api_key_dict=UserAPIKeyAuth(), + llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}), + ) + + assert _hidden_submit_model_warnings(caplog) == [] + + def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: return [ record.getMessage() diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 8da07f39fd8..7e0736d4e67 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1561,6 +1561,30 @@ async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert "applied_policies" not in data["metadata"] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + @pytest.mark.asyncio async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry From 00381ef03baedc9c9f7bec55f8c6139ecdb75f21 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 15:07:18 -0700 Subject: [PATCH 054/157] test: validate opaque stream IDs and hide log-reader credentials --- .../test_responses_bridge_streaming_e2e.py | 24 ++++++++++++------- tests/e2e/logging/datadog_reader.py | 10 ++++---- tests/e2e/logging/test_datadog_reader.py | 20 ++++++++++++++++ 3 files changed, 40 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/logging/test_datadog_reader.py diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py index 9a45743a0cd..75817340876 100644 --- a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p from __future__ import annotations +from typing import Final, Literal + import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import StreamingResponse @@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel): class _BridgeChunk(BaseModel): id: str - choices: list[_BridgeChoice] = [] + object: Literal["chat.completion.chunk"] + choices: list[_BridgeChoice] = Field(default_factory=list) class _WeatherArgs(BaseModel): @@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming: resources.key(), ChatBody( model=bridged_model, - messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], max_tokens=64, stream=True, ), ) - chunks = _bridge_chunks(result) - ids = {chunk.id for chunk in chunks} + chunks: Final = _bridge_chunks(result) + assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk" + ids: Final = frozenset(chunk.id for chunk in chunks) assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" - assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id" @pytest.mark.covers( "llm.chat_completions.openai.basic.stream.bridge_streams_sse", @@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming: chunks = _bridge_chunks(result) content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" - assert any( - choice.finish_reason for chunk in chunks for choice in chunk.choices - ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), ( + f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + ) assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" @pytest.mark.covers( diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index d0f478185c2..7b4372ef9ba 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -13,7 +13,7 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations import time -from dataclasses import dataclass +from dataclasses import dataclass, field import pytest from pydantic import BaseModel, ConfigDict, Field @@ -36,8 +36,8 @@ _RATE_LIMIT_RETRIES = 5 class _DdAuthHeaders(Headers): - api_key: str = Field(serialization_alias="DD-API-KEY") - app_key: str = Field(serialization_alias="DD-APPLICATION-KEY") + api_key: str = Field(serialization_alias="DD-API-KEY", repr=False) + app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False) class _SearchFilter(BaseModel): @@ -88,8 +88,8 @@ class _SearchResponse(BaseModel): @dataclass(frozen=True, slots=True) class DdLogsReader: site: str - api_key: str - app_key: str + api_key: str = field(repr=False) + app_key: str = field(repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py new file mode 100644 index 00000000000..00624bd3c84 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,20 @@ +from typing import Final + +from datadog_reader import DdLogsReader +from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization + + +def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: + api_key: Final = "test-datadog-api-secret" + app_key: Final = "test-datadog-app-secret" + reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key) + headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key) + + for value in (reader, headers): + assert api_key not in repr(value) + assert app_key not in repr(value) + + assert headers.model_dump(by_alias=True) == { + "DD-API-KEY": api_key, + "DD-APPLICATION-KEY": app_key, + } From 80782d478ddbdd2e145b0cb72eb3c653efbe0ed0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:17:36 -0700 Subject: [PATCH 055/157] fix(policy_engine): keep legacy stream steps in sync and reject tool-call rewrites Each streaming step drops the response an earlier step's translation stored under request_data["response"], so a later legacy hook sees the stream as the steps before it left it instead of the first step's snapshot. A legacy replacement whose tool calls differ from the scanned chunks is now undeliverable like a text mismatch, so the original stream is released with a warning instead of delivering the text while dropping the tool-call change --- .../proxy/policy_engine/pipeline_executor.py | 17 ++-- .../policy_engine/test_pipeline_executor.py | 78 ++++++++++++++++--- 2 files changed, 81 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 089f0683823..b4f6da18d1c 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -150,8 +150,8 @@ class _LegacyHookStreamAdapter(CustomGuardrail): non-streaming hooks, an exception it raises ends the stream through the executor's fail/error classification, and a replacement response is re-scanned by the same translation so its texts reach the client through the translation's ended-stream write-back. A - replacement whose scanned texts do not line up with the originals is undeliverable, so the - executor releases the original chunks.""" + replacement whose scanned texts do not line up with the originals, or whose tool calls + differ from them, is undeliverable, so the executor releases the original chunks.""" def __init__( self, @@ -184,8 +184,12 @@ class _LegacyHookStreamAdapter(CustomGuardrail): return inputs scanned: Final = _text_snapshot(inputs.get("texts")) rescanned: Final = await self._rescan(replacement, logging_obj) - rewritten: Final = None if rescanned is None else rescanned.get("texts") - if scanned is None or rewritten is None or len(rewritten) != len(scanned): + if scanned is None or rescanned is None: + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + rewritten: Final = rescanned.get("texts") + if rewritten is None or len(rewritten) != len(scanned): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} return rewritten_inputs @@ -380,7 +384,9 @@ class PipelineExecutor: yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the client gets - the stream the merge base sent.""" + the stream the merge base sent. The response an earlier step's translation stored under + ``request_data["response"]`` is dropped first, so this step's hook sees the stream as + the steps before it left it.""" scanner: Final = ( callback if PipelineExecutor.supports_unified_execution(callback) @@ -389,6 +395,7 @@ class PipelineExecutor: observer: Final = _StreamRewriteObserver(scanner) deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites originals: Final = copy.deepcopy(streaming_chunks) + hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored try: if deliver_rewrites: await endpoint_translation.process_output_streaming_response( diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index a44afc6ef96..860f83a06bb 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -675,7 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - result = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, @@ -1298,8 +1297,8 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri class _LegacyHookGuardrail(CustomGuardrail): """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" - def __init__(self, replacement=None, raises=None): - super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + def __init__(self, replacement=None, raises=None, guardrail_name="masker"): + super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) self.replacement = replacement self.raises = raises self.calls = [] @@ -1339,7 +1338,7 @@ class _LegacyScanningTranslation: ): request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) outputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [responses_so_far[0]["text"]]}, + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, @@ -1350,8 +1349,11 @@ class _LegacyScanningTranslation: async def process_output_response( self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None ): + inputs = {"texts": list(response["texts"])} + if response.get("tool_calls"): + inputs["tool_calls"] = list(response["tool_calls"]) await guardrail_to_apply.apply_guardrail( - inputs={"texts": list(response["texts"])}, + inputs=inputs, request_data={"response": response}, input_type="response", logging_obj=litellm_logging_obj, @@ -1359,10 +1361,26 @@ class _LegacyScanningTranslation: return response +def _legacy_replacement(*texts, tool_calls=None): + return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} + + async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): - monkeypatch.setattr(litellm, "callbacks", [guardrail]) + return await _run_legacy_streaming_steps(monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error) + + +async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail="block", on_error="next"): + monkeypatch.setattr(litellm, "callbacks", list(guardrails)) return await PipelineExecutor.execute_steps( - steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail=on_fail, on_error=on_error)], + steps=[ + PipelineStep( + guardrail=guardrail.guardrail_name, + on_pass="next" if position + 1 < len(guardrails) else "allow", + on_fail=on_fail, + on_error=on_error, + ) + for position, guardrail in enumerate(guardrails) + ], mode="post_call", data={"model": "m"}, user_api_key_dict=MagicMock(), @@ -1376,7 +1394,7 @@ async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="bl @pytest.mark.asyncio @pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): - guardrail = guardrail_class(replacement={"texts": ["[REWRITTEN] hello world"]}) + guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world")) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -1434,7 +1452,7 @@ async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatc @pytest.mark.asyncio async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): - guardrail = _LegacyHookGuardrail(replacement={"texts": ["split", "in two"]}) + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two")) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -1442,3 +1460,45 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail( + replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call]) + ) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[])) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): + masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) + auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor") + chunks = [_chunk()] + + result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next") + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass", "pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in masker.calls] == [{"native": True, "text": "hello world"}] + assert [call["response"] for call in auditor.calls] == [{"native": True, "text": "[REWRITTEN] hello world"}] From 64afa9d6eca8d238da02458a2769294eff93d08d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 15:36:34 -0700 Subject: [PATCH 056/157] test: isolate auto-router scenarios and clean partial setup --- .../test_auto_router_regressions_e2e.py | 239 +++++++++--------- 1 file changed, 120 insertions(+), 119 deletions(-) diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 188db2a8eb5..374badcf5fc 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form. import json import os from collections.abc import Iterator +from contextlib import ExitStack from dataclasses import dataclass from typing import Final @@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel): @dataclass(frozen=True, slots=True) -class TagSplitDeployments: - """Scenario A mirrors the customer-shaped config from GitHub issue #36619: - plain deployment registered first, tier deployment and marker both tagged. - Scenario B flips both axes for GitHub issue #36621: marker registered first - and its tier deployment left untagged, so routing depends neither on - registration order nor on tier deployments carrying tags.""" - - tag_a: str - shared_a: str - tier_a: str - tag_b: str - shared_b: str - tier_b: str +class TagSplitDeployment: + tag: str + shared: str + tier: str @dataclass(frozen=True, slots=True) @@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]: } -def _key_for( - proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False -) -> str: +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str: key: Final = proxy.generate_key( KeyGenerateBody( models=models, @@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con ) -@pytest.fixture(scope="module") -def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _register_models( + proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...] +) -> None: + for name, params in registrations: + stack.callback(proxy.delete_model, proxy.create_model(name, params)) + + +def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment: marker: Final = unique_marker() - deployments: Final = TagSplitDeployments( - tag_a=f"e2e-split-a-{marker}", - shared_a=f"e2e-autoroute-a-{marker}", - tier_a=f"e2e-tier-a-{marker}", - tag_b=f"e2e-split-b-{marker}", - shared_b=f"e2e-autoroute-b-{marker}", - tier_b=f"e2e-tier-b-{marker}", + named: Final = TagSplitDeployment( + tag=f"e2e-split-{marker}", + shared=f"e2e-autoroute-{marker}", + tier=f"e2e-tier-{marker}", ) anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") - marker_params_a: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_a), - tags=[deployments.tag_a], + marker_registration: Final = ( + named.shared, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + tags=[named.tag], + ), ) - marker_params_b: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_b), - tags=[deployments.tag_b], + tier_registration: Final = ( + named.tier, + LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]), ) - registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( - (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), - (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), - (deployments.shared_a, marker_params_a), - (deployments.shared_b, marker_params_b), - (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), - (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)) + registrations: Final = ( + (marker_registration, tier_registration, plain_registration) + if marker_first + else (plain_registration, tier_registration, marker_registration) ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield deployments - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, stack, registrations) + return named -@pytest.fixture(scope="module") -def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: +@pytest.fixture(scope="class") +def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=False) + + +@pytest.fixture(scope="class") +def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=True) + + +@pytest.fixture(scope="class") +def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias: marker: Final = unique_marker() named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: +@pytest.fixture(scope="class") +def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit: marker: Final = unique_marker() named: Final = HeuristicSplit( alias=f"e2e-heuristic-router-{marker}", @@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: +@pytest.fixture(scope="class") +def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter: marker: Final = unique_marker() named: Final = SemanticAutoRouter( marker=f"e2e-semantic-router-{marker}", @@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.marker, marker_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: +@pytest.fixture(scope="class") +def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias: marker: Final = unique_marker() named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named class TestTagSplitRouting: @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") def test_body_tagged_chat_routes_through_the_marker_to_its_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36619: with tag filtering on, a chat request whose body metadata tags match the tagged marker under a shared model name is answered by the marker's tier deployment, not by the plain deployment that was registered under the name first.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag]))) assert chat.choices, "tagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_chat_is_always_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620: untagged chat requests to the shared name succeed on every call and are all served by the plain deployment; the tagged marker never captures them, so no intermittent auto-router errors and no tier hijacking.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) for _ in range(5): - chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared))) assert chat.choices, "untagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=5) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + _assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_messages_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620 on the /v1/messages surface: an untagged Anthropic-native request to the shared name is served by the plain deployment, not captured by the tagged marker.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared))) assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name" + ) class TestUntaggedTierDeployments: @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36621: a /v1/messages request tagged only via the x-litellm-tags header selects the tagged marker, and the rewrite still lands on the tier deployment even though that deployment carries no tags, because the marker consumed the routing tags.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag) answer: Final = unwrap( proxy.transport.post( "/v1/messages", headers=headers, - json=_hello_messages_body(split.shared_b), + json=_hello_messages_body(marker_first_split.shared), response_type=AnthropicMessagesResponse, ) ) assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + _assert_served_only_by( + rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name" + ) @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins the tag-consumption half of GitHub issue #36621: after the tagged marker rewrites the request to its tier model, the consumed routing tags no longer constrain deployment selection, so the untagged tier deployment serves the request instead of a strict-tag denial.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + chat: Final = unwrap( + proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag])) + ) assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + _assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier") @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """The tag-consumption fix must not loosen strict tag semantics: a tagged request aimed directly at an untagged deployment (no marker involved) is still rejected with the 401 tags-configuration error.""" - key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) - result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag])) assert isinstance(result, UnauthorizedError), ( f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" ) @@ -451,37 +450,39 @@ class TestUntaggedTierDeployments: class TestResponsesApiTagRouting: @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_header_tagged_responses_with_string_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the /v1/responses surface of the tag split (GitHub issues #36620/#36621): a /v1/responses request with string input, tagged via the x-litellm-tags header, succeeds and routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) ) assert answer.id, "header-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + _assert_served_only_by( + rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input" + ) @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_body_tagged_responses_with_list_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the body-tag and list-input combination of the same split: /v1/responses with litellm_metadata.tags and structured input items routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, + model=plain_first_split.shared, input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], max_output_tokens=64, - litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]), ) answer: Final = unwrap( proxy.transport.post( @@ -493,18 +494,18 @@ class TestResponsesApiTagRouting: ) assert answer.id, "body-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_responses_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the untagged half of the /v1/responses tag split: an untagged request to the shared name is served by the plain deployment, matching the chat and messages surfaces.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post( @@ -516,7 +517,9 @@ class TestResponsesApiTagRouting: ) assert answer.id, "untagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name" + ) class TestStrategyAliasPricing: @@ -551,9 +554,7 @@ class TestComplexityHeuristicScope: while the accompanying ~2KB agent system prompt is packed with enough reasoning and complexity keywords that scoring the combined text lands in REASONING; only ask-only scoring keeps this on the cheap tier.""" - key: Final = _key_for( - proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] - ) + key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]) body: Final = ChatBody( model=heuristic_split.alias, messages=[ From 3ea65e761ee2f021e967c87202267cd4d70ff3f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:30:06 -0700 Subject: [PATCH 057/157] test(guardrails): cover Responses and Messages pipeline tool-call delivery --- .../proxy_logging/test_guardrail_pipeline.py | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index f753d2ab690..bf71780ad7b 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -25,6 +25,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines @@ -2185,3 +2186,106 @@ async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_str assert result is not None assert seen["count"] == 1 assert seen["response"] == "hello " + + +def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]: + return { + "tool_calls": [ + { + "id": stream_item_field(tool_call, "id"), + "type": "function", + "function": { + "name": stream_item_field(stream_item_field(tool_call, "function"), "name"), + "arguments": '{"fruit": "[MASKED]"}', + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + } + + +def _anthropic_tool_use_sse_chunks() -> List[bytes]: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw + assert "persim" not in raw + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert raw.count("event: content_block_delta") == 2 + + +def _responses_function_call_events() -> List[Dict[str, Any]]: + def item(arguments: str, status: str) -> Dict[str, Any]: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'}, + {"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'}, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"}, + }, + ] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()] + assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""] + assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert "persimmon" not in json.dumps(delivered) From ea427e33d8208034f934e52a7293f7f0bfa020ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:44:14 -0700 Subject: [PATCH 058/157] fix(policy_engine): deliver legacy stream rewrites made in place A legacy post-call hook that changes the response it was handed and returns None (the model armor guardrail masks choice content that way) used to have that rewrite dropped on the streaming pipeline path, since the adapter only re-scanned a returned replacement. The adapter now re-scans the response it handed the hook when the hook returns None, so an in-place rewrite reaches the client through the same ended-stream write-back --- .../proxy/policy_engine/pipeline_executor.py | 13 ++++--- .../policy_engine/test_pipeline_executor.py | 37 +++++++++++++++---- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index b4f6da18d1c..45a113fcda2 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -148,8 +148,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): endpoint translation hands it the texts it scanned plus the assembled response under ``request_data["response"]``; the hook gets that response in the shape its route gives non-streaming hooks, an exception it raises ends the stream through the executor's - fail/error classification, and a replacement response is re-scanned by the same translation - so its texts reach the client through the translation's ended-stream write-back. A + fail/error classification, and the response it hands back, or the one it changed in place + and returned ``None`` for, is re-scanned by the same translation so its texts reach the + client through the translation's ended-stream write-back. A replacement whose scanned texts do not line up with the originals, or whose tool calls differ from them, is undeliverable, so the executor releases the original chunks.""" @@ -175,15 +176,17 @@ class _LegacyHookStreamAdapter(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: + hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) replacement: Final = await self.inner.async_post_call_success_hook( data=request_data, user_api_key_dict=self.user_api_key_dict, - response=self.endpoint_translation.post_call_hook_response(request_data.get("response")), + response=hooked, ) - if replacement is None: + rewrite: Final = hooked if replacement is None else replacement + if rewrite is None: return inputs scanned: Final = _text_snapshot(inputs.get("texts")) - rescanned: Final = await self._rescan(replacement, logging_obj) + rescanned: Final = await self._rescan(rewrite, logging_obj) if scanned is None or rescanned is None: raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten: Final = rescanned.get("texts") diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 860f83a06bb..ae7754e5821 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1297,16 +1297,19 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri class _LegacyHookGuardrail(CustomGuardrail): """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" - def __init__(self, replacement=None, raises=None, guardrail_name="masker"): + def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None): super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) self.replacement = replacement self.raises = raises + self.rewrite_in_place = rewrite_in_place self.calls = [] async def async_post_call_success_hook(self, data, user_api_key_dict, response): self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) if self.raises is not None: raise self.raises + if self.rewrite_in_place is not None: + response["text"] = self.rewrite_in_place return self.replacement @@ -1325,7 +1328,7 @@ class _LegacyScanningTranslation: delivers_ended_stream_text_rewrites = True def post_call_hook_response(self, response): - return {"native": True, "text": response["text"]} + return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]} async def process_output_streaming_response( self, @@ -1336,7 +1339,9 @@ class _LegacyScanningTranslation: request_data=None, deliver_ended_stream_rewrites=False, ): - request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) + request_data.setdefault( + "response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]} + ) outputs = await guardrail_to_apply.apply_guardrail( inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, request_data=request_data, @@ -1349,7 +1354,7 @@ class _LegacyScanningTranslation: async def process_output_response( self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None ): - inputs = {"texts": list(response["texts"])} + inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])} if response.get("tool_calls"): inputs["tool_calls"] = list(response["tool_calls"]) await guardrail_to_apply.apply_guardrail( @@ -1361,6 +1366,10 @@ class _LegacyScanningTranslation: return response +def _native(text): + return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} + + def _legacy_replacement(*texts, tool_calls=None): return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} @@ -1403,12 +1412,26 @@ async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypa assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] assert chunks[0]["text"] == "[REWRITTEN] hello world" - assert [call["response"] for call in guardrail.calls] == [{"native": True, "text": "hello world"}] + assert [call["response"] for call in guardrail.calls] == [_native("hello world")] assert guardrail.calls[0]["data"]["model"] == "m" assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] assert not any("discarded" in record.getMessage() for record in caplog.records) +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world") + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert not any("discarded" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): guardrail = _LegacyHookGuardrail(replacement=None) @@ -1500,5 +1523,5 @@ async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(mon assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass", "pass"] assert chunks[0]["text"] == "[REWRITTEN] hello world" - assert [call["response"] for call in masker.calls] == [{"native": True, "text": "hello world"}] - assert [call["response"] for call in auditor.calls] == [{"native": True, "text": "[REWRITTEN] hello world"}] + assert [call["response"] for call in masker.calls] == [_native("hello world")] + assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")] From ce7cec1a3635bdb66292daa45ec9fe39f37bbffa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:48:22 -0700 Subject: [PATCH 059/157] fix(policy_engine): keep discarded stream rewrites out of the applied-guardrails header A streaming step whose rewrite the executor threw away (a tool-call rewrite, a text rewrite the translation cannot write back, or one the adapter refused) still marked its guardrail as applied, so the header claimed an output the client never received. The step now returns right after releasing the original chunks, which leaves the header as the merge base sent it --- litellm/proxy/policy_engine/pipeline_executor.py | 10 ++++++---- .../proxy/policy_engine/test_pipeline_executor.py | 1 + 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 45a113fcda2..9b0d1e60839 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -387,7 +387,8 @@ class PipelineExecutor: yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the client gets - the stream the merge base sent. The response an earlier step's translation stored under + the stream the merge base sent, and the guardrail stays out of the applied-guardrails + header since its output never reached the client. The response an earlier step's translation stored under ``request_data["response"]`` is dropped first, so this step's hook sees the stream as the steps before it left it.""" scanner: Final = ( @@ -419,9 +420,10 @@ class PipelineExecutor: ) except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) - else: - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + return + if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index ae7754e5821..2af318c88cd 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1146,6 +1146,7 @@ def _assert_passed_with_discard_warning(result, caplog): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @pytest.mark.asyncio From 253600fc6144747c1ffd954c3ebdfadd686adc2e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 16:52:56 -0700 Subject: [PATCH 060/157] test: wait for requested guardrail propagation --- tests/e2e/guardrails/guardrails_client.py | 22 ++++++- .../e2e/guardrails/test_guardrails_client.py | 63 +++++++++++++++++++ .../test_tool_permission_guardrail_e2e.py | 24 ++++--- 3 files changed, 98 insertions(+), 11 deletions(-) create mode 100644 tests/e2e/guardrails/test_guardrails_client.py diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 1f55a0f9a56..3770de0c6d5 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -7,7 +7,7 @@ from __future__ import annotations import time from collections.abc import Callable from dataclasses import dataclass -from typing import Literal +from typing import Final, Literal from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap @@ -405,6 +405,26 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient: return GuardrailsClient(proxy=proxy) +def poll_until_guardrail_applied( + call: Callable[[], StreamingResponse], + guardrail_name: str, + *, + timeout: float = POLL_TIMEOUT, + interval: float = POLL_INTERVAL, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> StreamingResponse: + deadline: Final = now() + timeout + while ( + (result := call()).ok + and guardrail_name + not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(",")) + and (remaining := deadline - now()) > 0 + ): + sleep(min(interval, remaining)) + return result + + def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]: """Retry a call that a guardrail should reject until it is, returning the last result. diff --git a/tests/e2e/guardrails/test_guardrails_client.py b/tests/e2e/guardrails/test_guardrails_client.py new file mode 100644 index 00000000000..aab163d86d2 --- /dev/null +++ b/tests/e2e/guardrails/test_guardrails_client.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass +from itertools import chain, repeat +from typing import Final + +import pytest + +from e2e_http import StreamingResponse +from guardrails_client import poll_until_guardrail_applied + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +def _response(applied: str, status: int = 200) -> StreamingResponse: + return StreamingResponse(status_code=status, body="{}", headers={"x-litellm-applied-guardrails": applied}) + + +def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> None: + clock: Final = Clock() + expected: Final = _response("global-filter, tool-permission") + responses: Final = iter((_response("global-filter"), expected)) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is expected + assert clock.elapsed == 2 + + +@pytest.mark.parametrize("applied", ("", "global-filter", "tool-permission-sibling")) +def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None: + clock: Final = Clock() + missing: Final = _response(applied) + + result: Final = poll_until_guardrail_applied( + lambda: missing, "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is missing + assert clock.elapsed == 5 + + +@pytest.mark.parametrize("status", (400, 401, 429, 500)) +def test_http_failure_is_not_hidden_by_a_later_success(status: int) -> None: + clock: Final = Clock() + failed: Final = _response("", status) + responses: Final = iter(chain((failed,), repeat(_response("tool-permission")))) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is failed + assert clock.elapsed == 0 diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py index 9ef3650625c..8d1047e53c7 100644 --- a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -30,6 +30,7 @@ from guardrails_client import ( ToolPermissionParamsBody, ToolPermissionRuleBody, poll_until_blocked, + poll_until_guardrail_applied, ) from lifecycle import ResourceManager from models import ChatResponse, ChatTool, ChatToolFunction @@ -84,8 +85,8 @@ def _register_tool_permission(client: GuardrailsClient, resources: ResourceManag resources.defer(lambda: client.delete_guardrail(guardrail_id)) -def _applied_guardrails(outcome: StreamingResponse) -> str: - return outcome.headers.get("x-litellm-applied-guardrails", "") +def _applied_guardrails(outcome: StreamingResponse) -> tuple[str, ...]: + return tuple(name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(",")) def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: @@ -144,14 +145,17 @@ class TestToolPermissionPreCall: name = f"e2e-toolperm-allow-{unique_marker()}" _register_tool_permission(client, resources, name=name) - outcome = client.chat_raw( - scoped_key, - MODEL, - TOOL_PROMPT, - guardrails=[name], - max_tokens=128, - tools=[ALLOWED_TOOL], - tool_choice="required", + outcome = poll_until_guardrail_applied( + lambda: client.chat_raw( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[ALLOWED_TOOL], + tool_choice="required", + ), + name, ) assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" From 9cb5d9b76c0c6b6e7a20e852e9673c9ea410c3f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:55:19 -0700 Subject: [PATCH 061/157] fix(proxy): gate streaming pipelines on a route-resolved guardrail translation The per-chunk hook skipped pipeline-managed guardrails whenever the request route was empty, while the gated stream could still fail to resolve a translation and release the buffered stream ungoverned. The gate now needs a translation resolved from the route, the iterator hook resolves it once and hands it to the gated stream, and the ungoverned release branch is gone. --- litellm/proxy/utils.py | 27 ++++------ .../proxy_logging/test_guardrail_pipeline.py | 53 +++++++++++++++---- 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c38eb900a05..37d077d607c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -192,6 +192,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -559,7 +560,7 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool: - return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None + return resolve_endpoint_translation(user_api_key_dict, None) is not None def _stream_gated_guardrail_names( @@ -3435,12 +3436,16 @@ class ProxyLogging: ), ) - if post_call_pipelines: + pipeline_translation: Final = ( + resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None + ) + if pipeline_translation is not None: current_response = self._pipeline_gated_stream( response=current_response, user_api_key_dict=user_api_key_dict, request_data=request_data, pipelines=post_call_pipelines, + translation=pipeline_translation, ) try: @@ -3464,6 +3469,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, request_data: dict, # mutable-ok: same request-payload shape the hooks mutate pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + translation: "tuple[str, BaseTranslation]", ) -> "AsyncGenerator[Any, None]": """ Execute post_call policy pipelines against a streamed response. @@ -3478,9 +3484,8 @@ class ProxyLogging: rewritten chunks, so rewrites chain). A rewrite the translation cannot deliver yet (one on a route without write-back, or a shape the route refuses) is discarded by the executor and the original chunks are - released, as is a buffered shape no translation resolves; a block or - modify_response terminates with the translation's block chunks or the - raised error. + released; a block or modify_response terminates with the translation's + block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3488,17 +3493,7 @@ class ProxyLogging: if not buffered: return - resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) - if resolved is None: - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; " - "the stream is released ungoverned by them: %s", - ", ".join(policy_name for policy_name, _pipeline in pipelines), - ) - for buffered_item in buffered: - yield buffered_item - return - call_type, endpoint_translation = resolved + call_type, endpoint_translation = translation for policy_name, pipeline in pipelines: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index bf71780ad7b..3b59bbac3fd 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1969,28 +1969,61 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe @pytest.mark.asyncio -async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape( +async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route( proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) - chunks = [object(), object()] - delivered: List[Any] = [] + chunks = _stream_chunks() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - async for item in proxy_logging.async_post_call_streaming_iterator_hook( - user_api_key_dict=make_user_api_key_auth(), - response=_async_chunk_iter(chunks), - request_data=data, - ): - delivered.append(item) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] assert len(delivered) == 2 assert seen.get("count") is None - assert any("response-governance" in message and "shape" in message for message in _warnings(caplog)) + assert any("response-governance" in message and "route None" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class UnifiedRecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + monkeypatch.setattr( + litellm, + "callbacks", + [UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(), + ) + + assert result is not None + assert seen["gr-post"] == 1 def _anthropic_sse_chunks() -> List[bytes]: From 88d2d775534457b3397c9332ca968b7a8707a161 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:59:55 -0700 Subject: [PATCH 062/157] feat(hosted_vllm): add image edit support Register a HostedVLLMImageEditConfig so hosted_vllm/ deployments route POST /v1/images/edits to the vLLM-Omni OpenAI-compatible endpoint instead of failing with 'image edit is not supported for hosted_vllm' before any request is sent --- .../llms/hosted_vllm/image_edit/__init__.py | 9 ++ .../hosted_vllm/image_edit/transformation.py | 42 +++++++ litellm/utils.py | 4 + ...t_hosted_vllm_image_edit_transformation.py | 114 ++++++++++++++++++ 4 files changed, 169 insertions(+) create mode 100644 litellm/llms/hosted_vllm/image_edit/__init__.py create mode 100644 litellm/llms/hosted_vllm/image_edit/transformation.py create mode 100644 tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py diff --git a/litellm/llms/hosted_vllm/image_edit/__init__.py b/litellm/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..27e005e8a0d --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .transformation import HostedVLLMImageEditConfig + +__all__ = ("HostedVLLMImageEditConfig",) + + +def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig: + return HostedVLLMImageEditConfig() diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py new file mode 100644 index 00000000000..bd27c962da8 --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -0,0 +1,42 @@ +"""Image edits for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/images/edits).""" + +from typing import Final + +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str + + +class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + """ + vLLM-Omni images edits API follows the OpenAI multipart contract. + + https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/images_api/ + """ + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseImageEditConfig contract + model: str, + api_key: str | None = None, + litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract + api_base: str | None = None, + ) -> dict: # mutable-ok: BaseImageEditConfig contract + resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" + return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseImageEditConfig contract + ) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError( + "api_base not set for Hosted VLLM images edits API. " + "Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable" + ) + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1"): + return f"{trimmed}/images/edits" + return f"{trimmed}/v1/images/edits" diff --git a/litellm/utils.py b/litellm/utils.py index 141ec323776..0f2b2f9042d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9239,6 +9239,10 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) + if LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config + + return get_hosted_vllm_image_edit_config(model=model) elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_edit.transformation import ( AzureImageEditConfig, diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py new file mode 100644 index 00000000000..e1167b92553 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -0,0 +1,114 @@ +"""Tests for hosted_vllm image edits (vLLM-Omni /v1/images/edits).""" + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config +from litellm.llms.hosted_vllm.image_edit.transformation import HostedVLLMImageEditConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng" +MODEL = "Qwen/Qwen-Image-Edit-2511" + + +@pytest.fixture(autouse=True) +def _clear_hosted_vllm_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + + +def test_provider_config_registration(): + config = ProviderConfigManager.get_provider_image_edit_config( + model=f"hosted_vllm/{MODEL}", + provider=LlmProviders.HOSTED_VLLM, + ) + + assert isinstance(config, HostedVLLMImageEditConfig) + assert isinstance(get_hosted_vllm_image_edit_config(MODEL), HostedVLLMImageEditConfig) + + +@pytest.mark.parametrize( + "api_base", + ["http://localhost:8091", "http://localhost:8091/", "http://localhost:8091/v1", "http://localhost:8091/v1/"], +) +def test_get_complete_url_appends_images_edits(api_base: str): + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=api_base, litellm_params={}) + == "http://localhost:8091/v1/images/edits" + ) + + +def test_get_complete_url_falls_back_to_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm-omni:8000/v1") + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + == "http://vllm-omni:8000/v1/images/edits" + ) + + +def test_get_complete_url_requires_api_base(): + config = HostedVLLMImageEditConfig() + + with pytest.raises(ValueError, match="api_base not set"): + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + + +def test_validate_environment_defaults_to_fake_api_key(): + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers == {"Authorization": "Bearer fake-api-key"} + + +def test_validate_environment_uses_provided_api_key_and_keeps_headers(): + headers = HostedVLLMImageEditConfig().validate_environment( + headers={"X-Test": "1"}, + model=MODEL, + api_key="my-custom-key", + ) + + assert headers == {"X-Test": "1", "Authorization": "Bearer my-custom-key"} + + +def test_validate_environment_falls_back_to_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "env-key") + + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers["Authorization"] == "Bearer env-key" + + +def test_image_edit_posts_multipart_to_vllm_omni(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + response = litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + api_key="test-key", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + seed=42, + ) + + assert response.data + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "http://localhost:8091/v1/images/edits" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'name="image[]"' in request.content + assert PNG_BYTES in request.content + assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content + assert b'name="prompt"\r\n\r\nadd a hat' in request.content + assert b'name="seed"\r\n\r\n42' in request.content From b8be30219c5e5dea4343b3a924b89494e7350303 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 17:02:00 -0700 Subject: [PATCH 063/157] test: stop guardrail retries at the polling deadline --- tests/e2e/guardrails/guardrails_client.py | 7 +++++-- tests/e2e/guardrails/test_guardrails_client.py | 5 ++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 3770de0c6d5..ed112a79b9b 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -415,13 +415,16 @@ def poll_until_guardrail_applied( sleep: Callable[[float], None] = time.sleep, ) -> StreamingResponse: deadline: Final = now() + timeout + if not (result := call()).ok: + return result while ( - (result := call()).ok - and guardrail_name + guardrail_name not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(",")) and (remaining := deadline - now()) > 0 ): sleep(min(interval, remaining)) + if now() >= deadline or not (result := call()).ok: + break return result diff --git a/tests/e2e/guardrails/test_guardrails_client.py b/tests/e2e/guardrails/test_guardrails_client.py index aab163d86d2..423c2ede599 100644 --- a/tests/e2e/guardrails/test_guardrails_client.py +++ b/tests/e2e/guardrails/test_guardrails_client.py @@ -40,13 +40,16 @@ def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None: clock: Final = Clock() missing: Final = _response(applied) + responses: Final = iter((missing, missing, missing)) result: Final = poll_until_guardrail_applied( - lambda: missing, "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep ) assert result is missing assert clock.elapsed == 5 + with pytest.raises(StopIteration): + next(responses) @pytest.mark.parametrize("status", (400, 401, 429, 500)) From 0c58346ba925d0a15ebcb38b162b9224446a50d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:07:37 -0700 Subject: [PATCH 064/157] fix(proxy): warn when a deferred background policy was matched through a request tag Retrieval re-matches only the key, team, and model scopes, so a post_call policy that reached a pending background response through a request-body tag does not govern the completed response. Log that at submit, next to the deferral, and cover the retrieval re-match with tag-scoped tests. --- litellm/proxy/utils.py | 22 +++++++++ .../policy_engine/test_response_retrieval.py | 22 +++++++++ .../proxy_logging/test_guardrail_pipeline.py | 47 ++++++++++++++++++- 3 files changed, 89 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4b62fc4cdae..d144aa113d6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -613,9 +613,31 @@ def _defer_post_call_pipelines( response.status, ", ".join(policy_name for policy_name, _pipeline in deferred), ) + tag_matched: Final = _tag_matched_deferrals(data, deferred) + if tag_matched: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: %s", + response.id, + ", ".join(f"{policy_name} ({source})" for policy_name, source in tag_matched), + ) _withdraw_deferred_claims(data, deferred) +def _tag_matched_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[tuple[str, str], ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + if not isinstance(sources, dict): + return () + return tuple( + (policy_name, str(sources[policy_name])) + for policy_name, _pipeline in deferred + if policy_name in sources and "tag:" in str(sources[policy_name]) + ) + + def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 9dcbef888ce..1fa1161191d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -68,6 +68,7 @@ def policy_engine(): "response-governance": _pipeline_policy("output-word-filter"), "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), "team-governance": _pipeline_policy("team-word-filter"), + "tag-governance": _pipeline_policy("tag-word-filter"), } ) attachment_registry.load_attachments( @@ -75,6 +76,7 @@ def policy_engine(): {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, {"policy": "team-governance", "teams": ["governed-team"]}, + {"policy": "tag-governance", "tags": ["governed"]}, ] ) yield @@ -119,6 +121,26 @@ def test_key_and_team_context_also_governs_retrieval(policy_engine): assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine): + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + + +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine): + data: dict[str, object] = { + "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), + "litellm_metadata": {"tags": ["governed"]}, + } + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),) + assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} + + def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): data = _retrieval_data(UNGOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index d6a74d3bbd9..63afce5d968 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1489,7 +1489,9 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] -def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None): +def _claimed_post_call_pipeline_data( + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str = "model:m" +): from litellm.proxy.policy_engine.policy_registry import get_policy_registry step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} @@ -1511,7 +1513,7 @@ def _claimed_post_call_pipeline_data(*policy_names: str, extra_guardrails: dict[ "_pipeline_managed_guardrails": {"gr-post"}, "applied_policies": list(policy_names), "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], - "policy_sources": {policy_name: "model:m" for policy_name in policy_names}, + "policy_sources": {policy_name: policy_source for policy_name in policy_names}, }, } @@ -1542,6 +1544,47 @@ async def test_pending_background_response_withdraws_the_deferred_policy_claims( assert "applied_guardrails" not in data["metadata"] +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert [ + message for message in _warnings(caplog) if "response-governance (tag:governed+model:m)" in message + ] == [ + "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: response-governance (tag:governed+model:m)" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_matched_through_its_model_does_not_warn( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog +): + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.post_call_success_hook( + data=_claimed_post_call_pipeline_data("response-governance"), + response=_background_response("queued"), + user_api_key_dict=make_user_api_key_auth(), + ) + + assert _warnings(caplog) == [] + + @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry From 831a2a13fbd87aee9e77cc6831963ac4a9975fbb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:12:18 -0700 Subject: [PATCH 065/157] fix(azure): price azure_ai transcriptions at the azure_ai cost-map entry --- litellm/llms/azure/audio_transcriptions.py | 7 ++- litellm/main.py | 1 + .../llms/azure/test_audio_transcriptions.py | 61 +++++++++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/azure/test_audio_transcriptions.py diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 4a5ed2ccb0c..564ec94ba6b 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -37,6 +37,7 @@ class AzureAudioTranscription(AzureChatCompletion): azure_ad_token: str | None = None, atranscription: bool = False, litellm_params: dict | None = None, + custom_llm_provider: str = "azure", ) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: data: Final = {"model": model, "file": audio_file, **optional_params} @@ -53,6 +54,7 @@ class AzureAudioTranscription(AzureChatCompletion): logging_obj=logging_obj, model=model, litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) azure_client: Final = self.get_azure_openai_client( @@ -99,7 +101,7 @@ class AzureAudioTranscription(AzureChatCompletion): additional_args={"complete_input_dict": data}, original_response=stringified_response, ) - hidden_params: Final = {"model": model, "custom_llm_provider": "azure"} + hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider} final_response: Final[TranscriptionResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, @@ -122,6 +124,7 @@ class AzureAudioTranscription(AzureChatCompletion): client=None, max_retries=None, litellm_params: dict | None = None, + custom_llm_provider: str = "azure", ) -> TranscriptionResponse: response = None try: @@ -178,7 +181,7 @@ class AzureAudioTranscription(AzureChatCompletion): }, original_response=stringified_response, ) - hidden_params: Final = {"model": model, "custom_llm_provider": "azure"} + hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider} response = convert_to_model_response_object( _response_headers=headers, response_object=stringified_response, diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..11b30c240aa 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7805,6 +7805,7 @@ def transcription( azure_ad_token=azure_ad_token, max_retries=max_retries, litellm_params=litellm_params_dict, + custom_llm_provider=custom_llm_provider, ) elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers): api_base = ( diff --git a/tests/test_litellm/llms/azure/test_audio_transcriptions.py b/tests/test_litellm/llms/azure/test_audio_transcriptions.py new file mode 100644 index 00000000000..cd5fcbd85a9 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_audio_transcriptions.py @@ -0,0 +1,61 @@ +import json +from pathlib import Path +from typing import Final + +import httpx +import pytest +from openai import AzureOpenAI + +import litellm +from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration + +AUDIO_FILE: Final = Path(__file__).parents[3] / "gettysburg.wav" +WHISPER_COST_PER_SECOND: Final = 0.0001 + + +def _transcription_client() -> AzureOpenAI: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"text": "Four score and seven years ago"}) + + return AzureOpenAI( + api_key="test-key", + api_version="2024-06-01", + azure_endpoint="https://example.cognitiveservices.azure.com", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def test_azure_ai_transcription_is_priced_at_the_azure_ai_entry(): + with AUDIO_FILE.open("rb") as audio: + response = litellm.transcription( + model="azure_ai/whisper", + file=audio, + api_base="https://example.cognitiveservices.azure.com", + api_key="test-key", + api_version="2024-06-01", + client=_transcription_client(), + ) + with AUDIO_FILE.open("rb") as audio: + duration = calculate_request_duration(audio) + + assert duration is not None and duration > 0 + assert response._hidden_params["custom_llm_provider"] == "azure_ai" + assert completion_cost(completion_response=response, call_type="transcription") == pytest.approx( + WHISPER_COST_PER_SECOND * duration + ) + + +def test_azure_transcription_keeps_the_azure_provider(): + with AUDIO_FILE.open("rb") as audio: + response = litellm.transcription( + model="azure/whisper-1", + file=audio, + api_base="https://example.openai.azure.com", + api_key="test-key", + api_version="2024-06-01", + client=_transcription_client(), + ) + + assert response._hidden_params["custom_llm_provider"] == "azure" + assert json.loads(response.model_dump_json())["text"] == "Four score and seven years ago" From 260097b2dde7f12f77784892e96e9956b0b19cc1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:16:34 -0700 Subject: [PATCH 066/157] refactor(hosted_vllm): drop redundant image edit docstrings --- litellm/llms/hosted_vllm/image_edit/transformation.py | 8 -------- .../test_hosted_vllm_image_edit_transformation.py | 2 -- 2 files changed, 10 deletions(-) diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index bd27c962da8..fa8c1dbfcfd 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -1,5 +1,3 @@ -"""Image edits for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/images/edits).""" - from typing import Final from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig @@ -7,12 +5,6 @@ from litellm.secret_managers.main import get_secret_str class HostedVLLMImageEditConfig(OpenAIImageEditConfig): - """ - vLLM-Omni images edits API follows the OpenAI multipart contract. - - https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/images_api/ - """ - def validate_environment( self, headers: dict, # mutable-ok: BaseImageEditConfig contract diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py index e1167b92553..5646f3d9d13 100644 --- a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -1,5 +1,3 @@ -"""Tests for hosted_vllm image edits (vLLM-Omni /v1/images/edits).""" - import httpx import pytest From 9e67c083092da5e20b85ed87e935d0af168cc66d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:16:39 -0700 Subject: [PATCH 067/157] fix(policy_engine): keep tool-only legacy streams deliverable A Messages stream that ends with only tool_use blocks reaches the legacy step without a texts key, while the non-streaming rescan of the same response sends an empty list. Compare both as empty and keep the stream when the hook left the tool calls alone. --- .../proxy/policy_engine/pipeline_executor.py | 15 +++- .../policy_engine/test_pipeline_executor.py | 84 ++++++++++++++++++- 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9b0d1e60839..52985d49c7e 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: return None if texts is None else tuple(texts) +def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]: + return tuple(texts or ()) + + def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) @@ -152,7 +156,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): and returned ``None`` for, is re-scanned by the same translation so its texts reach the client through the translation's ended-stream write-back. A replacement whose scanned texts do not line up with the originals, or whose tool calls - differ from them, is undeliverable, so the executor releases the original chunks.""" + differ from them, is undeliverable, so the executor releases the original chunks. A stream + that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as + long as the hook left the tool calls alone.""" def __init__( self, @@ -185,15 +191,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail): rewrite: Final = hooked if replacement is None else replacement if rewrite is None: return inputs - scanned: Final = _text_snapshot(inputs.get("texts")) rescanned: Final = await self._rescan(rewrite, logging_obj) - if scanned is None or rescanned is None: + if rescanned is None: raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten: Final = rescanned.get("texts") - if rewritten is None or len(rewritten) != len(scanned): + if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if not rewritten: + return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} return rewritten_inputs diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 2af318c88cd..98674843c73 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1367,6 +1367,44 @@ class _LegacyScanningTranslation: return response +class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation): + """Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from + the inputs, while the non-streaming scan of the same response sends an empty list.""" + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]}) + await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +def _tool_only_chunk(): + return {"text": "", "tool_call": _chunk()["tool_call"]} + + def _native(text): return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} @@ -1375,11 +1413,17 @@ def _legacy_replacement(*texts, tool_calls=None): return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} -async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): - return await _run_legacy_streaming_steps(monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error) +async def _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None +): + return await _run_legacy_streaming_steps( + monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation + ) -async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail="block", on_error="next"): +async def _run_legacy_streaming_steps( + monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None +): monkeypatch.setattr(litellm, "callbacks", list(guardrails)) return await PipelineExecutor.execute_steps( steps=[ @@ -1397,7 +1441,7 @@ async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail=" call_type="completion", policy_name="p", streaming_chunks=chunks, - endpoint_translation=_LegacyScanningTranslation(), + endpoint_translation=_LegacyScanningTranslation() if translation is None else translation, ) @@ -1513,6 +1557,38 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( assert chunks == [_chunk()] +@pytest.mark.asyncio +async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert chunks == [_tool_only_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_tool_only_chunk()] + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) From 359c26aa1d1403b0dc10c2d7a47b984d13c468c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:27:46 -0700 Subject: [PATCH 068/157] test(policy_engine): cover the legacy stream paths with no response and no rescan A translation that hands the hook no assembled response leaves the stream as it is, and a rewrite the translation cannot rescan is released as the original stream with a warning. --- .../policy_engine/test_pipeline_executor.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 98674843c73..6de7f7b593d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1589,6 +1589,68 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only assert chunks == [_tool_only_chunk()] +class _ResponselessLegacyScanningTranslation(_LegacyScanningTranslation): + """Like a handler that never stores the assembled response under request_data["response"].""" + + def post_call_hook_response(self, response): + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + +class _UnscannableRewriteTranslation(_LegacyScanningTranslation): + """Like the chat handler on a response whose choices are plain dicts: the non-streaming scan + never hands anything to the guardrail.""" + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + return response + + +@pytest.mark.asyncio +async def test_streaming_step_leaves_the_stream_alone_when_the_hook_gets_no_response(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail() + chunks = [_chunk()] + + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ResponselessLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert guardrail.calls[0]["response"] is None + assert chunks == [_chunk()] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]")) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) From 113350756595b7954d5f830f8a5454895de12526 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:46:55 -0700 Subject: [PATCH 069/157] refactor: type the buffered stream rewrite helpers without Any --- .../llms/anthropic/chat/guardrail_translation/handler.py | 8 ++++---- .../responses/test_openai_responses_guardrail_handler.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cbbccac17f3..6f966ae1a02 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,7 +13,7 @@ Pattern Overview: """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, MutableSequence, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat @@ -1262,7 +1262,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _write_ended_stream_text_rewrite( - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place rewritten_text: str, ) -> None: """Deliver an ended-stream guardrail text rewrite by rewriting the @@ -1284,7 +1284,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _write_ended_stream_tool_call_rewrites( cls, - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place *, pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], post_guardrail_tool_calls: tuple[_ToolCallShape, ...], @@ -1346,7 +1346,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _rewrite_ended_stream_events( - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place rewrite_event: _SSEEventRewriter, ) -> None: """Replace every buffered event ``rewrite_event`` returns a rewrite for, in diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index fa1e969b277..6ed4ec6618f 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -17,6 +17,7 @@ from fastapi import HTTPException from openai.types.responses import ResponseFunctionToolCall from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1238,7 +1239,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: tool_calls = [ {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} From 24ee66328a7bf1a791dc0f7d81051165611877c9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 17:55:21 -0700 Subject: [PATCH 070/157] test: bound Datadog read-back retries using reset headers --- tests/e2e/logging/datadog_reader.py | 114 ++++++++----- tests/e2e/logging/test_datadog_reader.py | 203 +++++++++++++++++++++++ 2 files changed, 278 insertions(+), 39 deletions(-) diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index 7b4372ef9ba..368c20cb6aa 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations +import math +import random import time +from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,12 +31,28 @@ from e2e_config import ( DD_SITE, POLL_TIMEOUT, ) -from e2e_http import URL, Headers, RateLimitedError, Success, post +from e2e_http import URL, Headers, StreamingResponse, send -#: How many rate-limited responses in a row one search tolerates before the -#: hard fail; each retry sleeps a full search interval, so this rides out a -#: burst from a concurrent consumer of the org-wide search budget. -_RATE_LIMIT_RETRIES = 5 +type SearchCall = Callable[[str, float], StreamingResponse] + + +def _seconds(value: str | None) -> float | None: + if value is None: + return None + try: + seconds: Final = float(value) + except ValueError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + +def _rate_limit_delay(headers: Mapping[str, str]) -> float: + delays: Final = tuple( + delay + for name in ("x-ratelimit-reset", "retry-after") + if (delay := _seconds(headers.get(name))) is not None + ) + return max(1.0, max(delays, default=DD_SEARCH_INTERVAL)) class _DdAuthHeaders(Headers): @@ -90,6 +110,10 @@ class DdLogsReader: site: str api_key: str = field(repr=False) app_key: str = field(repr=False) + search: SearchCall | None = field(default=None, repr=False) + now: Callable[[], float] = field(default=time.monotonic, repr=False) + sleep: Callable[[float], None] = field(default=time.sleep, repr=False) + jitter: Callable[[], float] = field(default=random.random, repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog @@ -108,25 +132,28 @@ class DdLogsReader: a single event. A 429 backs off and retries - the search budget is org-wide, so another consumer can empty it under us - while any other failure stays a hard fail.""" - for _ in range(_RATE_LIMIT_RETRIES): - result = post( - URL(f"https://api.{self.site}/api/v2/logs/events/search"), - headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=query)), - response_type=_SearchResponse, - timeout=30.0, - ) - match result: - case Success(data=page): - return [event.attributes for event in page.data] - case RateLimitedError(retry_after_seconds=retry_after): - time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL) - case failure: - pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}") + return self._events_for_query(query, self.now() + POLL_TIMEOUT) + + def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]: + search: Final = self.search or self._search_page + while (remaining := deadline - self.now()) > 0: + if (result := search(query, min(30.0, remaining))).ok: + return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data] + if result.status_code != 429: + pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}") + if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0: + self.sleep(delay) pytest.fail( - f"DataDog Logs Search API at api.{self.site} still rate-limited after " - f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide " - "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer" + f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; " + "the org-wide logs_public_search_api budget is exhausted" + ) + + def _search_page(self, query: str, timeout: float) -> StreamingResponse: + return send( + URL(f"https://api.{self.site}/api/v2/logs/events/search"), + headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), + json=_SearchRequest(filter=_SearchFilter(query=query)), + timeout=timeout, ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: @@ -140,33 +167,42 @@ class DdLogsReader: hide from the exactly-one assertion - real-DataDog jitter can surface one call's two events tens of seconds apart. Searches pace at DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's - request budget. At the deadline the last result is returned as-is.""" - deadline = time.monotonic() + POLL_TIMEOUT - while time.monotonic() < deadline: - events = self.events_for_query(query) + request budget. Discovery, quota retries, and duplicate detection share + one POLL_TIMEOUT deadline; an incomplete settle window fails closed.""" + deadline: Final = self.now() + POLL_TIMEOUT + while (remaining := deadline - self.now()) > 0: + events = self._events_for_query(query, deadline) if events: - return self._settled_events_for_query(query, events) - time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_query(query) + return self._settled_events_for_query(query, events, deadline) + if (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + return [] - def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. Keep the last non-empty result: a transient empty search (index lag) must not erase events already confirmed earlier in the settle window. + A successful final search must reach the full settle window before the + shared read-back deadline; otherwise duplicate detection is incomplete. """ - settle_deadline = time.monotonic() + DD_SETTLE_SECONDS + settle_deadline: Final = self.now() + DD_SETTLE_SECONDS last_nonempty = events - while time.monotonic() < settle_deadline: - time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_query(query) - if not latest: - continue + if len(events) > 1: + return events + while (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + if self.now() >= deadline: + break + latest = self._events_for_query(query, deadline) if len(latest) > 1: return latest - last_nonempty = latest - return last_nonempty + if latest: + last_nonempty = latest + if self.now() >= settle_deadline: + return last_nonempty + pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s") def build_dd_logs_reader() -> DdLogsReader: diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py index 00624bd3c84..910a1cefd42 100644 --- a/tests/e2e/logging/test_datadog_reader.py +++ b/tests/e2e/logging/test_datadog_reader.py @@ -1,7 +1,14 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass from typing import Final +import pytest + from datadog_reader import DdLogsReader from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization +from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT +from e2e_http import StreamingResponse def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: @@ -18,3 +25,199 @@ def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> "DD-API-KEY": api_key, "DD-APPLICATION-KEY": app_key, } + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +@dataclass +class Search: + responses: Iterator[StreamingResponse] + calls: tuple[tuple[str, float], ...] = () + + def __call__(self, query: str, timeout: float) -> StreamingResponse: + self.calls += ((query, timeout),) + return next(self.responses) + + +def _page(*event_ids: str) -> StreamingResponse: + return StreamingResponse( + status_code=200, + body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}), + ) + + +def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]: + search: Final = Search(iter(responses)) + return DdLogsReader( + site="us5.datadoghq.com", + api_key="test-api-secret", + app_key="test-app-secret", + search=search, + now=clock.now, + sleep=clock.sleep, + jitter=lambda: 0.25, + ), search + + +def test_429_honors_server_reset_and_preserves_duplicate_events() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")), + clock, + ) + + events: Final = reader.events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 6.25 + assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0)) + + +@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1")) +def test_invalid_reset_uses_search_interval(reset: str) -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25 + + +def test_zero_reset_cannot_create_a_busy_retry_loop() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 1.25 + + +def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 8.25 + + +def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls == (("test-marker", 30.0),) + + +def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75)) + + +@pytest.mark.parametrize("status", (-1, 401, 403, 500)) +def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None: + clock: Final = Clock() + reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock) + + with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"): + reader.events_for_query("test-marker") + + assert search.calls == (("test-marker", 30.0),) + assert clock.elapsed == 0 + + +def test_polling_quota_retries_share_the_original_deadline() -> None: + clock: Final = Clock() + reader, search = _reader( + (_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == 2 + + +def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None: + clock: Final = Clock() + attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) + reader, search = _reader((_page(),) * attempts, clock) + + assert reader.poll_events_for_query("test-marker") == [] + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == attempts + + +def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader( + (_page(),) * empty_reads + + (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL) + assert len(search.calls) == empty_reads + 2 + + +def test_settlement_detects_a_duplicate_on_the_final_search() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 30 + + +def test_settlement_keeps_confirmed_events_through_empty_searches() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first",) + assert clock.elapsed == 30 + + +def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock) + + with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == empty_reads + 2 From 163847333242b7e75b1a15c27519beef2b1365ed Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 19:02:08 -0700 Subject: [PATCH 071/157] fix(bedrock): keep deletion response IDs in request context --- litellm/llms/bedrock/files/transformation.py | 18 ++++++------- .../test_bedrock_files_transformation.py | 27 +++++++++++++++++++ 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 90b539ff37c..9875ac2b9c3 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -13,7 +13,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_logger @@ -61,7 +61,11 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" -S3_DELETE_FILE_ID_PARAM: Final = "_s3_delete_file_id" + + +class _S3DeleteContext(BaseModel): + file_id: str = Field(min_length=1) + # litellm_params key carrying the size of the body uploaded to S3, handed from # `transform_create_file_request` to `transform_create_file_response`. @@ -1187,11 +1191,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: MutableMapping[str, object], ) -> tuple[str, dict[str, str]]: - request: Final = self._transform_s3_file_request( + return self._transform_s3_file_request( file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params ) - litellm_params[S3_DELETE_FILE_ID_PARAM] = file_id - return request def transform_delete_file_response( self, @@ -1205,10 +1207,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}", headers=raw_response.headers, ) - file_id: Final = litellm_params.get(S3_DELETE_FILE_ID_PARAM) - if not isinstance(file_id, str) or not file_id: - raise ValueError("Missing file id for Bedrock file deletion response") - return FileDeleted(id=file_id, deleted=True, object="file") + context: Final = _S3DeleteContext.model_validate(logging_obj.model_call_details.get("additional_args")) + return FileDeleted(id=context.file_id, deleted=True, object="file") def transform_list_files_request( self, diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 2c02a58663e..3b01a4f2054 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1861,6 +1861,33 @@ class TestBedrockFileDeletion: S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl" URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl" + def test_interleaved_deletions_keep_their_own_file_ids(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + config: Final = BedrockFilesConfig() + params: Final = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + } + file_ids: Final = (self.S3_URI, "s3://my-bucket/litellm-bedrock-files-model-second.jsonl") + for file_id in file_ids: + config.transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params=params) + + deleted: Final = tuple( + config.transform_delete_file_response( + raw_response=httpx.Response(204), + logging_obj=MagicMock(model_call_details={"additional_args": {"file_id": file_id}}), + litellm_params=params, + ).id + for file_id in file_ids + ) + + assert deleted == file_ids + def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None: import httpx import respx From 0086b62b4575fb5bb69653c2ad3d196822bd3f13 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 02:05:08 +0000 Subject: [PATCH 072/157] fix(cost-map): keep first fetch blocking, run retries in background Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 4 +- .../litellm_core_utils/get_model_cost_map.py | 132 +++++++++- litellm/proxy/proxy_server.py | 23 +- .../test_get_model_cost_map.py | 235 +++++++++++++----- 4 files changed, 296 insertions(+), 98 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..4e3754399da 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -541,7 +541,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount @@ -2397,3 +2397,5 @@ def __getattr__(name: str) -> Any: # ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time + +mark_litellm_import_complete() diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 9cba5db8ab7..4118cd3420e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -12,6 +12,7 @@ import asyncio import json import os import random +import threading import time from collections.abc import Awaitable, Callable from dataclasses import dataclass @@ -159,6 +160,15 @@ class GetModelCostMap: RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0 +_litellm_import_complete = threading.Event() + + +def mark_litellm_import_complete() -> None: + _litellm_import_complete.set() + + +def _start_daemon_thread(fn: Callable[[], None]) -> None: + threading.Thread(target=fn, name="litellm-model-cost-map-retry", daemon=True).start() @dataclass(frozen=True, slots=True) @@ -297,9 +307,15 @@ def _fetch_remote_model_cost_map_with_retry_sync( sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, + starting_attempt: int = 1, + initial_outcome: _FetchAttemptRetryable | None = None, ) -> ModelCostMapReloadResult: - for attempt in range(1, max_attempts + 1): - outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) + for attempt in range(starting_attempt, max_attempts + 1): + outcome = ( + initial_outcome + if initial_outcome is not None and attempt == starting_attempt + else _attempt_fetch_sync(client=client, url=url, timeout=timeout) + ) if not isinstance(outcome, _FetchAttemptRetryable): return outcome wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) @@ -464,6 +480,70 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) +def adopt_model_cost_map( + new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract +) -> int: + import litellm + from litellm import utils + + litellm.model_cost = new_model_cost_map + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + litellm.add_known_models(model_cost_map=new_model_cost_map) + fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 + utils.reapply_runtime_model_cost_registrations() + return fetched_model_count + + +def _continue_remote_fetch_in_background( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, + first_outcome: _FetchAttemptRetryable, + apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict +) -> None: + try: + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=rng, + client=client, + initial_outcome=first_outcome, + ) + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", + url, + max_attempts, + ) + return + backup_model_count: Final = GetModelCostMap._get_backup_model_count() # pyright: ignore[reportPrivateUsage] # integrity cache + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=backup_model_count, + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Keeping local backup. url=%s", + url, + ) + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + return + finalized_map: Final = _finalize_model_cost_map(result.model_cost_map) + _litellm_import_complete.wait() + apply(finalized_map) + _cost_map_source_info.source = "remote" + _cost_map_source_info.url = url + _cost_map_source_info.is_env_forced = False + _cost_map_source_info.fallback_reason = None + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + except Exception as e: + verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) + + def get_model_cost_map( url: str, timeout: int = 5, @@ -471,14 +551,19 @@ def get_model_cost_map( sleep: Callable[[float], None] = time.sleep, rng: random.Random | None = None, client: "_SyncGetClient | None" = None, + start_background: Callable[[Callable[[], None]], None] = _start_daemon_thread, + apply: Callable[ # mutable-ok: injected callback receives the mutable cost-map dict + [dict], + object, + ] = adopt_model_cost_map, ) -> dict: """ Public entry point — returns the model cost map dict. 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. - 2. Otherwise fetches from ``url``, retrying transient HTTP errors - (429/5xx/transport) with Retry-After-aware backoff, validates - integrity, and falls back to the local backup on any failure. + 2. Otherwise fetches from ``url``, validates the first response, and falls + back to the local backup while retrying transient HTTP errors in the + background. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -497,14 +582,35 @@ def get_model_cost_map( _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - result: Final = _fetch_remote_model_cost_map_with_retry_sync( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=rng if rng is not None else random.Random(), - client=client if client is not None else httpx, - ) + fetch_client: Final = client if client is not None else httpx + fetch_rng: Final = rng if rng is not None else random.Random() + first_outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) + if isinstance(first_outcome, _FetchAttemptRetryable): + verbose_logger.warning( + "LiteLLM: model cost map fetch attempt 1/%d failed: %s; " + "using local backup while retrying in the background", + max_attempts, + first_outcome.reason, + ) + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {first_outcome.reason}" + local_map: Final = _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + if max_attempts > 1: + start_background( + lambda: _continue_remote_fetch_in_background( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=fetch_rng, + client=fetch_client, + first_outcome=first_outcome, + apply=apply, + ) + ) + return local_map + + result: Final = first_outcome if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..efe38376353 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -132,11 +132,7 @@ from litellm.types.utils import ( TextCompletionResponse, TokenCountResponse, ) -from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - load_credentials_from_list, - reapply_runtime_model_cost_registrations, -) +from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from aiohttp import ClientSession @@ -4411,20 +4407,9 @@ def resolve_classifier_plugin( def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: - """Adopt a freshly fetched cost map into this process's litellm state, return the model count""" - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # Repopulate provider model sets (e.g. litellm.anthropic_models) so that - # wildcard patterns like "anthropic/*" include any newly added models. - litellm.add_known_models(model_cost_map=new_model_cost_map) - # Counted before the re-apply below, which writes into this same dict, so the - # number reported describes the fetched price data alone. - fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 - # The swap discards everything registered at runtime (deployment model_info, - # register_model overrides), so put it back on top of the fresh catalog. - reapply_runtime_model_cost_registrations() - return fetched_model_count + from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map + + return adopt_model_cost_map(new_model_cost_map) def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool: diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 18185126775..dce1a431f13 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -20,21 +20,18 @@ from litellm.litellm_core_utils.get_model_cost_map import ( GetModelCostMap, _count_model_entries, _finalize_model_cost_map, + adopt_model_cost_map, ) def _load_root_cost_map() -> dict: - path = os.path.join( - os.path.dirname(__file__), "../../../model_prices_and_context_window.json" - ) + path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json") with open(path) as f: return json.load(f) def _make_models(n: int) -> dict: - return { - f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n) - } + return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)} def test_count_model_entries_excludes_reserved_keys(): @@ -117,9 +114,7 @@ def test_finalize_pops_key_and_installs_rules(): def test_finalize_with_no_block_clears_rules(): previous = list(get_fallback_generalization_rules()) try: - set_fallback_generalizations( - [{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}] - ) + set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]) _finalize_model_cost_map(_make_models(2)) assert match_capability_generalizations("x-1") is None finally: @@ -306,9 +301,7 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch): from litellm.litellm_core_utils import get_model_cost_map as module monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) - client, _calls = _mock_client( - [httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client - ) + client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) before = datetime.now(timezone.utc) module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) @@ -317,6 +310,7 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch): assert loaded_at is not None assert before <= loaded_at <= datetime.now(timezone.utc) + # --------------------------------------------------------------------------- # refetch_model_cost_map: retry/backoff behavior for runtime reloads # --------------------------------------------------------------------------- @@ -382,9 +376,7 @@ async def test_refetch_retries_429_honoring_retry_after(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 assert calls["count"] == 3 @@ -396,9 +388,7 @@ async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(): """All 429 without Retry-After: exponential backoff waits, then a failure value.""" client, calls = _mock_client([httpx.Response(429)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "429" in result.reason assert "after 3 attempts" in result.reason @@ -418,9 +408,7 @@ async def test_refetch_caps_retry_after_wait(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert sleeper.waits == [30.0] @@ -435,9 +423,7 @@ async def test_refetch_retries_transport_errors(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert calls["count"] == 2 assert len(sleeper.waits) == 1 @@ -448,9 +434,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): """A 404 is permanent: one attempt, no sleeps, failure value.""" client, calls = _mock_client([httpx.Response(404)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "404" in result.reason assert calls["count"] == 1 @@ -461,9 +445,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): async def test_refetch_invalid_json_fails_immediately(): client, calls = _mock_client([httpx.Response(200, content=b"not json")]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "invalid JSON" in result.reason assert calls["count"] == 1 @@ -475,9 +457,7 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): """A drastically shrunk upstream file is rejected instead of being adopted.""" tiny = json.dumps(_make_models(60)).encode() client, _calls = _mock_client([httpx.Response(200, content=tiny)]) - result = await refetch_model_cost_map( - url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "integrity validation" in result.reason @@ -520,62 +500,187 @@ class _SyncSleepRecorder: self.waits.append(seconds) -def test_boot_load_retries_transient_failures_instead_of_falling_back(): - """A refused connection then a 503 at pod boot used to pin the process to the bundled - backup for its lifetime; both are transient and must be retried before giving up.""" +class _BackgroundRecorder: + def __init__(self): + self.callbacks = [] + + def __call__(self, callback): + self.callbacks.append(callback) + + +def test_boot_load_returns_local_map_and_schedules_transient_retry(): + client, calls = _mock_client( + [ + httpx.ConnectError("connection refused"), + ], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + + cost_map = get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert len(background.callbacks) == 1 + assert len(cost_map) > 100 + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None + + +def test_background_retry_adopts_valid_remote_map(): client, calls = _mock_client( [ httpx.ConnectError("connection refused"), - httpx.Response(503), httpx.Response(200, content=_real_map_bytes()), ], client_cls=httpx.Client, ) sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + applied = [] - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, + apply=applied.append, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert len(background.callbacks) == 1 + + background.callbacks[0]() + + assert calls["count"] == 2 + assert len(sleeper.waits) == 1 + assert 2.0 <= sleeper.waits[0] < 3.0 + assert len(applied) == 1 + assert applied[0].keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + + +def test_background_retry_keeps_local_map_after_remaining_failures(): + client, calls = _mock_client( + [httpx.ConnectError("connection refused"), httpx.Response(503)], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + applied = [] + + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, + apply=applied.append, + ) + background.callbacks[0]() assert calls["count"] == 3 assert len(sleeper.waits) == 2 assert 2.0 <= sleeper.waits[0] < 3.0 assert 4.0 <= sleeper.waits[1] < 5.0 + assert applied == [] + assert get_model_cost_map_source_info()["source"] == "local" + + +def test_boot_load_does_not_schedule_non_retryable_failure(): + client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, + ) + + assert calls["count"] == 1 + assert sleeper.waits == [] + assert background.callbacks == [] source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None + assert source["source"] == "local" + assert source["fallback_reason"] is not None + + +def test_boot_load_success_does_not_schedule_background_retry(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + + cost_map = get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert background.callbacks == [] + assert get_model_cost_map_source_info()["source"] == "remote" assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} -def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts(): - """An outage longer than the retry budget still ends on the bundled backup, and the - recorded fallback reason says how many attempts were spent so operators can tell.""" - client, calls = _mock_client( - [httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client +def test_boot_load_with_one_attempt_does_not_schedule_background_retry(): + client, calls = _mock_client([httpx.ConnectError("connection refused")], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + background = _BackgroundRecorder() + + get_model_cost_map( + url=_URL, + max_attempts=1, + sleep=sleeper, + rng=random.Random(0), + client=client, + start_background=background, ) - sleeper = _SyncSleepRecorder() - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - - assert calls["count"] == 3 - assert sleeper.waits == [7.0, 7.0] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert "after 3 attempts" in source["fallback_reason"] - assert len(cost_map) > 100 - - -def test_boot_load_does_not_retry_permanent_failures(): - """A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup.""" - client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - - get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert calls["count"] == 1 assert sleeper.waits == [] + assert background.callbacks == [] assert get_model_cost_map_source_info()["source"] == "local" - get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0)) - assert sleeper.waits == [] - assert get_model_cost_map_source_info()["source"] == "local" + +def test_adopt_model_cost_map_replays_runtime_registration_and_provider_models(): + import litellm + from litellm import utils as litellm_utils + + original_model_cost = litellm.model_cost + original_registry = dict(litellm_utils._runtime_registered_model_cost) + original_anthropic_models = set(litellm.anthropic_models) + try: + litellm.register_model( + model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} + ) + + models_count = adopt_model_cost_map({"anthropic/new-model": {"litellm_provider": "anthropic", "mode": "chat"}}) + + assert models_count == 1 + assert "anthropic/new-model" in litellm.anthropic_models + assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 + finally: + litellm.model_cost = original_model_cost # test-quality-ok: restore the module state changed by adoption + litellm_utils._runtime_registered_model_cost.clear() + litellm_utils._runtime_registered_model_cost.update(original_registry) + litellm.anthropic_models.clear() + litellm.anthropic_models.update(original_anthropic_models) + litellm_utils._invalidate_model_cost_lowercase_map() def test_boot_load_respects_local_env_override(monkeypatch): From ddedb4867b8cedea4ac2223dbc636f813e8bf997 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:09:44 -0700 Subject: [PATCH 073/157] fix: discard a streamed rewrite that drops or adds a tool call A guardrail that removes or adds a tool call on an ended stream used to be silently ignored: every handler substitutes the original list on a count mismatch and the executor skipped its observer once the translation could deliver rewrites. The executor now tracks the count change on the observer and releases the original chunks with the discard warning on every translation, matching what the merge base did for any tool call rewrite --- .../proxy/policy_engine/pipeline_executor.py | 20 ++++-- .../policy_engine/test_pipeline_executor.py | 23 +++++- .../proxy_logging/test_guardrail_pipeline.py | 71 +++++++++++++++++++ 3 files changed, 108 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c870bd7d2ec..a51468cc0fb 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -78,6 +78,10 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent +def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and len(returned) != len(sent) + + _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -91,8 +95,9 @@ class _StreamRewriteObserver(CustomGuardrail): guardrail. It records whether the guardrail returned different output than it was given, which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and tool-call rewrites are deliverable on translations that write them back across the - buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation - are discarded by the executor, which releases the original chunks. + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation, + and a rewrite that drops or adds a tool call on any translation, are discarded by the + executor, which releases the original chunks. The inner guardrail's ``apply_guardrail`` already records the guardrail information and span, so the observer's stays out of ``log_guardrail_information``.""" @@ -101,6 +106,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False + self.changed_tool_call_count = False def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -118,9 +124,11 @@ class _StreamRewriteObserver(CustomGuardrail): outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) + returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) - self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote( - sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) + self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + sent_tool_shapes, returned_tool_shapes ) return outputs @@ -325,7 +333,9 @@ class PipelineExecutor: except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) else: - if not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls): + if observer.changed_tool_call_count or ( + not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) + ): _release_original_chunks(step.guardrail, streaming_chunks, originals) if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 5d042d3c1ad..16401ccdaad 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1106,7 +1106,8 @@ class _WritingTranslation: logging_obj=litellm_logging_obj, ) responses_so_far[0]["text"] = outputs["texts"][0] - responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] + if len(outputs["tool_calls"]) == 1: + responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] return responses_so_far @@ -1242,6 +1243,26 @@ async def test_streaming_step_delivers_tool_call_rewrite_through_writing_transla assert not any("discarded" in record.getMessage() for record in caplog.records) +class _ToolCallDroppingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + @pytest.mark.asyncio async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 3b59bbac3fd..d46881d045b 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -2322,3 +2322,74 @@ async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_o assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' assert "persimmon" not in json.dumps(delivered) + + +def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]: + return {"tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_tool_call_stream_chunks()), + request_data=data, + ) + ] + + assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + assert delivered == _anthropic_tool_use_sse_chunks() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert delivered == _responses_function_call_events() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) From aa0a9ab3ea4dff53a22cbc60fbc0195c4ab6098c Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 02:29:56 +0000 Subject: [PATCH 074/157] refactor(cost-map): drop initial_outcome flag from retry loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 90ff696f926..d9b5a492539 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -324,19 +324,14 @@ async def _fetch_remote_model_cost_map_with_retry( def _fetch_remote_model_cost_map_with_retry_sync( url: str, timeout: int, - max_attempts: int, + attempts: range, sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, - starting_attempt: int = 1, - initial_outcome: _FetchAttemptRetryable | None = None, ) -> ModelCostMapReloadResult: - for attempt in range(starting_attempt, max_attempts + 1): - outcome = ( - initial_outcome - if initial_outcome is not None and attempt == starting_attempt - else _attempt_fetch_sync(client=client, url=url, timeout=timeout) - ) + max_attempts: Final = attempts.stop - 1 + for attempt in attempts: + outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) if not isinstance(outcome, _FetchAttemptRetryable): return outcome wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) @@ -557,18 +552,18 @@ def _continue_remote_fetch_in_background( sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, - first_outcome: _FetchAttemptRetryable, + first_wait: float, apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict ) -> None: try: + sleep(first_wait) result: Final = _fetch_remote_model_cost_map_with_retry_sync( url=url, timeout=timeout, - max_attempts=max_attempts, + attempts=range(2, max_attempts + 1), sleep=sleep, rng=rng, client=client, - initial_outcome=first_outcome, ) if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( @@ -652,19 +647,26 @@ def get_model_cost_map( local_map: Final = _finalize_loaded_model_cost_map( GetModelCostMap.load_local_model_cost_map_with_revision() ).model_cost_map - if max_attempts > 1: - start_background( - lambda: _continue_remote_fetch_in_background( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=fetch_rng, - client=fetch_client, - first_outcome=first_outcome, - apply=apply, - ) + first_wait: Final = _next_retry_wait( + outcome=first_outcome, + attempt=1, + max_attempts=max_attempts, + rng=fetch_rng, + ) + if isinstance(first_wait, ModelCostMapReloadUnavailable): + return local_map + start_background( + lambda: _continue_remote_fetch_in_background( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=fetch_rng, + client=fetch_client, + first_wait=first_wait, + apply=apply, ) + ) return local_map result: Final = first_outcome From 536a85b42967fd0c9c6b49d2df9298232034427a Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 02:48:17 +0000 Subject: [PATCH 075/157] fix(cost-map): keep register_model url fetch to a single attempt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 2 +- tests/test_litellm/test_utils.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 36b48d3b8d8..8df28870544 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3079,7 +3079,7 @@ def register_model( # Convert stringified numbers to appropriate numeric types loaded_model_cost = model_cost elif isinstance(model_cost, str): - loaded_model_cost = litellm.get_model_cost_map(url=model_cost) + loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1) if persist_across_reloads: _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index fee5e3a2e4c..e90372141bd 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2,10 +2,12 @@ import asyncio import json import logging import os +import threading from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from jsonschema import validate @@ -2382,6 +2384,27 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() +@respx.mock +def test_register_model_url_fetch_uses_single_attempt(monkeypatch): + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + before = dict(litellm.model_cost) + threads_before = {thread.name for thread in threading.enumerate()} + route = respx.get("https://example.invalid/custom_pricing.json").mock( + return_value=httpx.Response(503) + ) + + litellm.register_model(model_cost="https://example.invalid/custom_pricing.json") + + threads_after = {thread.name for thread in threading.enumerate()} + assert route.call_count == 1 + assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"} + assert not any( + thread.name == "litellm-model-cost-map-retry" and thread.is_alive() + for thread in threading.enumerate() + ) + assert litellm.model_cost.keys() >= before.keys() + + def test_register_model_openrouter_without_slash(): """ Test that register_model handles openrouter models without '/' in the name. From 4179f086e0196107dfc83da2a4e0bb1d15568edd Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 02:58:51 +0000 Subject: [PATCH 076/157] refactor(cost-map): share local-fallback and remote-accept paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 111 ++++++++---------- 1 file changed, 47 insertions(+), 64 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index d9b5a492539..67d06a0758e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -531,6 +531,32 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) +def _use_local_backup(reason: str | None) -> dict: # mutable-ok: returns the mutable model-cost map contract + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = reason + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + + +def _accept_remote( + result: ModelCostMapReloaded, url: str +) -> dict | None: # mutable-ok: returns the mutable model-cost map contract + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + return None + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _cost_map_source_info.source = "remote" + _cost_map_source_info.url = url + _cost_map_source_info.is_env_forced = False + _cost_map_source_info.fallback_reason = None + return finalized + + def adopt_model_cost_map( new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract ) -> int: @@ -545,17 +571,20 @@ def adopt_model_cost_map( return fetched_model_count -def _continue_remote_fetch_in_background( +def _retry_remote_fetch_in_background( url: str, timeout: int, max_attempts: int, sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, - first_wait: float, + first_outcome: _FetchAttemptRetryable, apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict ) -> None: try: + first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) + if isinstance(first_wait, ModelCostMapReloadUnavailable): + return sleep(first_wait) result: Final = _fetch_remote_model_cost_map_with_retry_sync( url=url, @@ -572,23 +601,12 @@ def _continue_remote_fetch_in_background( max_attempts, ) return - backup_model_count: Final = GetModelCostMap._get_backup_model_count() # pyright: ignore[reportPrivateUsage] # integrity cache - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=backup_model_count, - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Keeping local backup. url=%s", - url, - ) + _litellm_import_complete.wait() + accepted: Final = _accept_remote(result, url) + if accepted is None: _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" return - _litellm_import_complete.wait() - apply(_finalize_loaded_model_cost_map(result).model_cost_map) - _cost_map_source_info.source = "remote" - _cost_map_source_info.url = url - _cost_map_source_info.is_env_forced = False - _cost_map_source_info.fallback_reason = None + apply(accepted) _cost_map_source_info.loaded_at = datetime.now(timezone.utc) except Exception as e: verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) @@ -623,77 +641,42 @@ def get_model_cost_map( # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": - _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + return _use_local_backup(None) _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False fetch_client: Final = client if client is not None else httpx fetch_rng: Final = rng if rng is not None else random.Random() - first_outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) - if isinstance(first_outcome, _FetchAttemptRetryable): + outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) + if isinstance(outcome, ModelCostMapReloaded): + accepted: Final = _accept_remote(outcome, url) + return accepted if accepted is not None else _use_local_backup("Remote data failed integrity validation") + if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: verbose_logger.warning( "LiteLLM: model cost map fetch attempt 1/%d failed: %s; " "using local backup while retrying in the background", max_attempts, - first_outcome.reason, + outcome.reason, ) - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {first_outcome.reason}" - local_map: Final = _finalize_loaded_model_cost_map( - GetModelCostMap.load_local_model_cost_map_with_revision() - ).model_cost_map - first_wait: Final = _next_retry_wait( - outcome=first_outcome, - attempt=1, - max_attempts=max_attempts, - rng=fetch_rng, - ) - if isinstance(first_wait, ModelCostMapReloadUnavailable): - return local_map start_background( - lambda: _continue_remote_fetch_in_background( + lambda: _retry_remote_fetch_in_background( url=url, timeout=timeout, max_attempts=max_attempts, sleep=sleep, rng=fetch_rng, client=fetch_client, - first_wait=first_wait, + first_outcome=outcome, apply=apply, ) ) - return local_map - - result: Final = first_outcome - if isinstance(result, ModelCostMapReloadUnavailable): + else: verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - result.reason, + outcome.reason, ) - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - content: Final = result.model_cost_map - - # Validate using cached count (cheap int comparison, no file I/O) - if not GetModelCostMap.validate_model_cost_map( - fetched_map=content, - backup_model_count=GetModelCostMap._get_backup_model_count(), - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, - ) - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - - _cost_map_source_info.source = "remote" - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(result).model_cost_map + return _use_local_backup(f"Remote fetch failed: {outcome.reason}") From 94f9230d13900897f764498169bdfa3c5fa2a4ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:59:09 -0700 Subject: [PATCH 077/157] fix(proxy): type the new pipeline tests and keep tag values out of the deferral warning Every test this PR adds now annotates its fixture and parametrize parameters. The submit-time warning for a tag-matched deferred policy names only the policies, since a wildcard attachment pattern would let caller-provided tag text reach the log. --- litellm/proxy/utils.py | 6 +- .../policy_engine/test_response_retrieval.py | 42 ++-- .../proxy_logging/test_guardrail_pipeline.py | 216 ++++++++++-------- 3 files changed, 155 insertions(+), 109 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d144aa113d6..1604dea1003 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -620,19 +620,19 @@ def _defer_post_call_pipelines( "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " "does not govern the completed response: %s", response.id, - ", ".join(f"{policy_name} ({source})" for policy_name, source in tag_matched), + ", ".join(tag_matched), ) _withdraw_deferred_claims(data, deferred) def _tag_matched_deferrals( data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] -) -> tuple[tuple[str, str], ...]: +) -> tuple[str, ...]: sources: Final = _policy_state_metadata(data).get("policy_sources") if not isinstance(sources, dict): return () return tuple( - (policy_name, str(sources[policy_name])) + policy_name for policy_name, _pipeline in deferred if policy_name in sources and "tag:" in str(sources[policy_name]) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index 1fa1161191d..37849101b3a 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -1,5 +1,5 @@ import logging -from collections.abc import Mapping +from collections.abc import Iterator, Mapping import pytest @@ -60,7 +60,7 @@ def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, objec @pytest.fixture -def policy_engine(): +def policy_engine() -> Iterator[None]: policy_registry = get_policy_registry() attachment_registry = get_attachment_registry() policy_registry.load_policies( @@ -97,7 +97,7 @@ def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], .. ) -def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -111,7 +111,7 @@ def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine): assert "guardrails" not in data["litellm_metadata"] -def test_key_and_team_context_also_governs_retrieval(policy_engine): +def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None: data = _retrieval_data(UNGOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval( @@ -121,7 +121,7 @@ def test_key_and_team_context_also_governs_retrieval(policy_engine): assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) -def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine): +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -129,7 +129,7 @@ def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) -def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine): +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None: data: dict[str, object] = { "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), "litellm_metadata": {"tags": ["governed"]}, @@ -141,7 +141,7 @@ def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag( assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} -def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None: data = _retrieval_data(UNGOVERNED_MODEL_ID) attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) @@ -149,7 +149,7 @@ def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine): assert data == _retrieval_data(UNGOVERNED_MODEL_ID) -def test_already_attached_policy_is_not_attached_twice(policy_engine): +def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) router = _router() attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) @@ -168,7 +168,9 @@ def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str] ] -def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(policy_engine, caplog): +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: data = _retrieval_data(WILDCARD_MODEL_ID) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -176,11 +178,14 @@ def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( assert data == _retrieval_data(WILDCARD_MODEL_ID) assert [ - "as model group openai/* (a wildcard deployment)" in message for message in _hidden_submit_model_warnings(caplog) + "as model group openai/* (a wildcard deployment)" in message + for message in _hidden_submit_model_warnings(caplog) ] == [True] -def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_engine, caplog): +def test_aliased_model_group_still_attaches_its_own_policies_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) @@ -194,7 +199,9 @@ def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_en ] == [True] -def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(policy_engine, caplog): +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): attach_post_call_pipelines_to_retrieval( data=_retrieval_data(GOVERNED_MODEL_ID), @@ -209,7 +216,8 @@ def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str return [ record.getMessage() for record in caplog.records - if record.levelno == logging.WARNING and "retrieved without its post_call policy pipelines" in record.getMessage() + if record.levelno == logging.WARNING + and "retrieved without its post_call policy pipelines" in record.getMessage() ] @@ -221,7 +229,9 @@ def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str (None, "response id names no deployment"), ], ) -def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, caplog, response_id, reason): +def test_unresolvable_response_id_attaches_nothing_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str +) -> None: data = {"response_id": response_id, "litellm_metadata": {}} with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -231,7 +241,7 @@ def test_unresolvable_response_id_attaches_nothing_and_warns(policy_engine, capl assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): +def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None: data = _retrieval_data(GOVERNED_MODEL_ID) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -241,7 +251,7 @@ def test_without_a_router_attaches_nothing_and_warns(policy_engine, caplog): assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] -def test_without_policy_engine_attaches_nothing_quietly(caplog): +def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None: get_policy_registry().clear() data = _retrieval_data(GOVERNED_MODEL_ID) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 63afce5d968..3bbdbbe2ccd 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio import json import logging +from collections.abc import Iterator from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -155,9 +156,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_routes_through_router( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": cb}) @@ -173,9 +172,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_router_none_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth): with patch("litellm.proxy.proxy_server.llm_router", None): with pytest.raises(ValueError, match="Router not initialized"): await proxy_logging._execute_guardrail_with_load_balancing( @@ -188,9 +185,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_no_callback_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth): router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": None}) with patch("litellm.proxy.proxy_server.llm_router", router): @@ -210,9 +205,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises( @pytest.mark.asyncio -async def test_process_guardrail_callback_skipped_when_should_run_false( - proxy_logging, make_user_api_key_auth -): +async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=False) out = await proxy_logging._process_guardrail_callback( @@ -226,9 +219,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false( @pytest.mark.asyncio -async def test_process_guardrail_callback_returns_data_on_success( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=True) proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) @@ -343,14 +334,14 @@ async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging, @pytest.mark.asyncio -async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode( + proxy_logging, make_user_api_key_auth, monkeypatch +): pipeline = MagicMock() pipeline.mode = "post_call" # not pre_call data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []} executed = MagicMock() - monkeypatch.setattr( - "litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed - ) + monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed) out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), @@ -538,9 +529,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result( - result=result, data={"model": "m"}, policy_name="p" - ) + ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") finally: litellm.callbacks = saved @@ -652,9 +641,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics( - callback=cb, coro=task(), hook_type="post_call" - ) + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs @@ -682,9 +669,7 @@ def _moderation_guardrail() -> MagicMock: @pytest.mark.asyncio -async def test_during_call_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -703,9 +688,7 @@ async def test_during_call_hook_records_latency_metric( @pytest.mark.asyncio -async def test_post_call_success_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -733,9 +716,7 @@ async def test_post_call_success_hook_records_latency_metric( async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( data=data, @@ -760,9 +741,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -810,9 +789,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) with pytest.raises(RuntimeError): @@ -913,9 +890,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -1124,9 +1099,7 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1289,11 +1262,7 @@ async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes( monkeypatch.setattr( litellm, "callbacks", - [ - BlockingWriterGuardrail( - guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False - ) - ], + [BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data() @@ -1379,9 +1348,7 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1433,14 +1400,20 @@ def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail] seen["response"] = response raise HTTPException(status_code=400, detail={"error": "output blocked"}) - return [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + return [ + OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] @pytest.mark.asyncio @pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog, pending_status -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + pending_status: str, +) -> None: seen: dict[str, object] = {} monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1464,8 +1437,11 @@ async def test_post_call_success_hook_waits_for_pending_background_response_befo @pytest.mark.asyncio @pytest.mark.parametrize("final_status", ["completed", "incomplete"]) async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( - proxy_logging, make_user_api_key_auth, monkeypatch, final_status -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + final_status: str, +) -> None: seen: dict[str, object] = {} monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1486,7 +1462,9 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: async def async_post_call_success_hook(self, data, user_api_key_dict, response): return response - return [OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)] + return [ + OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] def _claimed_post_call_pipeline_data( @@ -1519,7 +1497,7 @@ def _claimed_post_call_pipeline_data( @pytest.fixture -def clear_policy_registry(): +def clear_policy_registry() -> Iterator[None]: from litellm.proxy.policy_engine.policy_registry import get_policy_registry yield @@ -1528,8 +1506,11 @@ def clear_policy_registry(): @pytest.mark.asyncio async def test_pending_background_response_withdraws_the_deferred_policy_claims( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance") @@ -1546,8 +1527,12 @@ async def test_pending_background_response_withdraws_the_deferred_policy_claims( @pytest.mark.asyncio async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") @@ -1559,19 +1544,21 @@ async def test_pending_background_response_warns_when_the_deferred_policy_was_ma assert out.status == "queued" assert "policy_sources" not in data["metadata"] - assert [ - message for message in _warnings(caplog) if "response-governance (tag:governed+model:m)" in message - ] == [ + assert [message for message in _warnings(caplog) if "through a request tag" in message] == [ "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " - "does not govern the completed response: response-governance (tag:governed+model:m)" + "does not govern the completed response: response-governance" ] @pytest.mark.asyncio async def test_pending_background_response_matched_through_its_model_does_not_warn( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) @@ -1587,12 +1574,17 @@ async def test_pending_background_response_matched_through_its_model_does_not_wa @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data( - "input-and-output-governance", "response-governance", extra_guardrails={"input-and-output-governance": ["gr-pre"]} + "input-and-output-governance", + "response-governance", + extra_guardrails={"input-and-output-governance": ["gr-pre"]}, ) await proxy_logging.post_call_success_hook( @@ -1606,8 +1598,11 @@ async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs @pytest.mark.asyncio async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: class DualStageGuardrail(CustomGuardrail): async def async_post_call_success_hook(self, data, user_api_key_dict, response): return response @@ -1630,8 +1625,11 @@ async def test_pending_background_response_keeps_the_claim_of_a_default_on_guard @pytest.mark.asyncio async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( - proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _claimed_post_call_pipeline_data("response-governance") @@ -1647,8 +1645,11 @@ async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pip @pytest.mark.asyncio async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", []) data = _post_call_pipeline_data(background=True) @@ -1709,7 +1710,9 @@ def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], ) pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) - data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} + data = { + "metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]} + } with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) @@ -2009,7 +2012,9 @@ def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): return {**inputs, **transform(inputs)} - return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + return RewritingStreamGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) def _tool_call_stream_chunks() -> List[Any]: @@ -2188,11 +2193,38 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvab def _anthropic_sse_chunks() -> List[bytes]: events = [ - ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}), + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), ("message_stop", {"type": "message_stop"}), ] return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] @@ -2259,7 +2291,13 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop assert "hello [MASKED]" in raw assert "hello world" not in raw assert raw.count("event: content_block_delta") == 1 - for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + for expected_event in ( + "message_start", + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + ): assert f"event: {expected_event}" in raw @@ -2332,9 +2370,7 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( managed = UnifiedRecordingGuardrail( guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True ) - free = RecordingGuardrail( - guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True - ) + free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True) monkeypatch.setattr(litellm, "callbacks", [managed, free]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) From c66ae07e3bb97e875d88096458e5de80b351828c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:11:01 -0700 Subject: [PATCH 078/157] fix(hosted_vllm): reject image edit params vLLM-Omni ignores --- .../hosted_vllm/image_edit/transformation.py | 9 ++++ litellm/utils.py | 2 +- ...t_hosted_vllm_image_edit_transformation.py | 43 +++++++++++++++++++ 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index fa8c1dbfcfd..3b8cc437168 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -3,8 +3,17 @@ from typing import Final from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str +PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"}) + class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract + return [ # mutable-ok: BaseImageEditConfig returns list + param + for param in super().get_supported_openai_params(model) + if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT + ] + def validate_environment( self, headers: dict, # mutable-ok: BaseImageEditConfig contract diff --git a/litellm/utils.py b/litellm/utils.py index 0f2b2f9042d..fc530f14e28 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9239,7 +9239,7 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) - if LlmProviders.HOSTED_VLLM == provider: + elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config return get_hosted_vllm_image_edit_config(model=model) diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py index 5646f3d9d13..bc0ea23e249 100644 --- a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -110,3 +110,46 @@ def test_image_edit_posts_multipart_to_vllm_omni(): assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content assert b'name="prompt"\r\n\r\nadd a hat' in request.content assert b'name="seed"\r\n\r\n42' in request.content + + +@pytest.mark.parametrize("param", ["mask", "quality", "input_fidelity"]) +def test_params_vllm_omni_ignores_are_not_advertised(param: str): + supported = HostedVLLMImageEditConfig().get_supported_openai_params(MODEL) + + assert param not in supported + assert {"image", "prompt", "n", "size", "response_format", "background", "user"} <= set(supported) + + +def test_image_edit_rejects_quality_unless_dropped(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(litellm.UnsupportedParamsError, match="quality"): + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + ) + assert captured == [] + + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + drop_params=True, + ) + + assert len(captured) == 1 + assert b'name="quality"' not in captured[0].content + assert b'name="prompt"\r\n\r\nadd a hat' in captured[0].content From 8d040d89e68c075b105a6369e5d95fdc7c1ba8f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:39:56 -0700 Subject: [PATCH 079/157] fix(policy_engine): keep legacy hooks off streams their route cannot assemble and off guardrails with their own iterator hook The streaming pipeline step only takes a post-call hook on routes whose translation assembles the streamed response (chat completions, Responses, Messages). On /v1/completions, the Gemini streamGenerateContent route, and A2A streams the pipeline is skipped with the merge-base warning and the hook runs on its own afterwards, instead of getting a None response while the header says the guardrail ran. A guardrail that overrides async_post_call_streaming_iterator_hook next to its post-call hook keeps its native per-chunk path rather than running buffered through the adapter --- .../chat/guardrail_translation/handler.py | 1 + .../guardrail_translation/base_translation.py | 7 ++ .../chat/guardrail_translation/handler.py | 1 + .../guardrail_translation/handler.py | 1 + .../proxy/policy_engine/pipeline_executor.py | 14 ++-- litellm/proxy/utils.py | 60 +++++++++++---- .../policy_engine/test_pipeline_executor.py | 46 +++-------- .../proxy_logging/test_guardrail_pipeline.py | 77 ++++++++++++++++++- 8 files changed, 148 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index f61cfe58e80..b83235f660f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -171,6 +171,7 @@ class AnthropicMessagesHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def __init__(self): super().__init__() diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index bef472b2882..375b22ff5f6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -60,6 +60,13 @@ class BaseTranslation(ABC): text rewrites on every other translation, are undeliverable: the pipeline executor discards them and releases the original chunks.""" + assembles_streamed_response: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` stores the assembled response of an + ended stream under ``request_data["response"]`` before scanning it, the way the chat, + Responses, and Messages translations do. A streaming pipeline runs a guardrail that only + has the legacy post-call hook against that response, so on a translation without it such + a guardrail keeps running on its own.""" + def post_call_hook_response(self, response: object) -> object: """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from the object the translation stores under ``request_data["response"]`` while scanning an diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 80292aef2cf..b9b88061fec 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -79,6 +79,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b0f79552bc5..e0ba80706df 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -341,6 +341,7 @@ class OpenAIResponsesHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 52985d49c7e..ab5ad01fda9 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -556,12 +556,14 @@ class PipelineExecutor: @staticmethod def supports_streaming_execution(callback: CustomGuardrail) -> bool: """Whether a streaming pipeline step can run this guardrail against the buffered - stream: through the unified path, or through its own post-call hook on the - assembled response. A guardrail with neither (one that only rewrites the stream - through its iterator hook) has to keep running on its own.""" - return ( - PipelineExecutor.supports_unified_execution(callback) - or type(callback).async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + stream: through the unified path, or through its post-call hook on the assembled + response when that hook is its only streaming path. A guardrail with its own + streaming iterator hook, or with neither hook, keeps running on its own.""" + callback_type: Final = type(callback) + return PipelineExecutor.supports_unified_execution(callback) or ( + callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + and callback_type.async_post_call_streaming_iterator_hook + is CustomLogger.async_post_call_streaming_iterator_hook ) @staticmethod diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 23260785b97..67a579c4d15 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -192,6 +192,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -517,9 +518,17 @@ def _merge_pipeline_metadata_writes( _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - return callback is not None and PipelineExecutor.supports_streaming_execution(callback) + if callback is None: + return False + if PipelineExecutor.supports_unified_execution(callback): + return True + return ( + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -541,42 +550,57 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No ) -def _pipeline_unsupported_streaming_guardrails(pipeline: "GuardrailPipeline") -> tuple[str, ...]: +def _pipeline_unsupported_streaming_guardrails( + pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> tuple[str, ...]: return tuple( dict.fromkeys( - step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) + step.guardrail + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail, translation) ) ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline) +def _pipeline_is_streamable( + policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation) if not unsupported: return True verbose_proxy_logger.warning( - "Policy '%s' has post_call pipeline guardrails with neither the unified apply_guardrail interface nor a " - "post-call hook, one of which streaming pipelines need; the stream skips the pipeline and its guardrails " - "run on their own: %s", + "Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they " + "need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a " + "route whose translation assembles the streamed response. The stream skips the pipeline and its " + "guardrails run on their own: %s", policy_name, ", ".join(unsupported), ) return False -def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool: - return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None +def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None": + resolved: Final = resolve_endpoint_translation(user_api_key_dict, None) + return None if resolved is None else resolved[1] + + +def _route_supports_streaming_pipelines( + user_api_key_dict: UserAPIKeyAuth, translation: "BaseTranslation | None" +) -> bool: + return not user_api_key_dict.request_route or translation is not None def stream_gated_guardrail_names( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: - if not _route_supports_streaming_pipelines(user_api_key_dict): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if not _route_supports_streaming_pipelines(user_api_key_dict, translation): return frozenset() return _pipeline_step_guardrail_names( tuple( (policy_name, pipeline) for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) ) ) @@ -589,8 +613,9 @@ def _streamable_post_call_pipelines( Streaming pipelines scan the buffered stream through the endpoint guardrail translation of the request route, so every step's guardrail needs either the - unified apply_guardrail interface or a post-call hook to run against the - assembled response, and the route needs a translation. A pipeline that + unified apply_guardrail interface or, on a route whose translation assembles + the streamed response, a post-call hook that is its only streaming path, and + the route needs a translation. A pipeline that cannot be run that way yet is left out and its guardrails run on the stream on their own, the way they did before pipelines ran on streams at all, with a warning naming the pipeline. @@ -598,7 +623,8 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - if not _route_supports_streaming_pipelines(user_api_key_dict): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if not _route_supports_streaming_pipelines(user_api_key_dict, translation): verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " @@ -610,7 +636,7 @@ def _streamable_post_call_pipelines( return tuple( (policy_name, pipeline) for policy_name, pipeline in post_call_pipelines - if _pipeline_is_streamable(policy_name, pipeline) + if _pipeline_is_streamable(policy_name, pipeline, translation) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 6de7f7b593d..e2523567065 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1589,28 +1589,14 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only assert chunks == [_tool_only_chunk()] -class _ResponselessLegacyScanningTranslation(_LegacyScanningTranslation): - """Like a handler that never stores the assembled response under request_data["response"].""" +class _NoHooksGuardrail(CustomGuardrail): + pass - def post_call_hook_response(self, response): - return response - async def process_output_streaming_response( - self, - responses_so_far, - guardrail_to_apply, - litellm_logging_obj=None, - user_api_key_dict=None, - request_data=None, - deliver_ended_stream_rewrites=False, - ): - await guardrail_to_apply.apply_guardrail( - inputs={"texts": [responses_so_far[0]["text"]]}, - request_data=request_data, - input_type="response", - logging_obj=litellm_logging_obj, - ) - return responses_so_far +class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item class _UnscannableRewriteTranslation(_LegacyScanningTranslation): @@ -1623,21 +1609,11 @@ class _UnscannableRewriteTranslation(_LegacyScanningTranslation): return response -@pytest.mark.asyncio -async def test_streaming_step_leaves_the_stream_alone_when_the_hook_gets_no_response(monkeypatch, caplog): - guardrail = _LegacyHookGuardrail() - chunks = [_chunk()] - - result = await _run_legacy_streaming_step( - monkeypatch, guardrail, chunks, translation=_ResponselessLegacyScanningTranslation() - ) - - assert result.terminal_action == "allow" - assert [step.outcome for step in result.step_results] == ["pass"] - assert guardrail.calls[0]["response"] is None - assert chunks == [_chunk()] - assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] - assert not any("discarded" in record.getMessage() for record in caplog.records) +def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path(): + assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False + assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 46d6f2bb7a9..ac303e8a027 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -28,7 +28,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header -from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines +from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( @@ -1539,6 +1539,21 @@ def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuar return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) +def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorAndLegacyHookGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1 + return None + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + def _rewritten_model_response(response: Any) -> litellm.ModelResponse: payload = response.model_dump() payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"] @@ -1572,6 +1587,42 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) +@pytest.mark.parametrize( + "request_route", + [None, "/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], +) +def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( + make_user_api_key_auth, monkeypatch, caplog, request_route +): + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})]) + legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}} + auth = make_user_api_key_auth(request_route=request_route) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})]) + both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}} + auth = make_user_api_key_auth(request_route="/v1/chat/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog)) + + def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( make_user_api_key_auth, monkeypatch, caplog ): @@ -1762,6 +1813,30 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen == {"iterator_hook_calls": 1} + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "rewrite_attribute, value", From 9a721abf0d098f40caa6ba67d343e99de74f7fc2 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 03:50:49 +0000 Subject: [PATCH 080/157] test(cost-map): clear LITELLM_LOCAL_MODEL_COST_MAP in register_model url test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e90372141bd..e42608c9904 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2386,6 +2386,7 @@ def test_register_model_with_scientific_notation(): @respx.mock def test_register_model_url_fetch_uses_single_attempt(monkeypatch): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) before = dict(litellm.model_cost) threads_before = {thread.name for thread in threading.enumerate()} From 15392e7b3a298bec57f4aae31e863e39fb182cc6 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 21:26:04 -0700 Subject: [PATCH 081/157] fix(mcp): surface connection test failures safely --- litellm/experimental_mcp_client/client.py | 21 +++ .../mcp_server/rest_endpoints.py | 59 ++++++- .../test_mcp_client.py | 124 ++++++++++++++- .../mcp_server/test_rest_endpoints.py | 148 +++++++++++++++++- 4 files changed, 334 insertions(+), 18 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3503468c735..cabcc6c03ba 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -18,6 +18,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.shared.message import SessionMessage +from mcp.shared.session import RequestResponder from typing_extensions import Unpack _TransportStreams: TypeAlias = tuple[ @@ -56,10 +57,13 @@ def missing_streamable_http_client_error() -> ImportError: from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( + ClientResult, GetPromptRequestParams, GetPromptResult, Prompt, ResourceTemplate, + ServerNotification, + ServerRequest, TextContent, ) from mcp.types import Tool as MCPTool @@ -442,6 +446,18 @@ class MCPClient: in_flight_error: BaseException | None = None try: read_stream, write_stream = transport[0], transport[1] + stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() + + async def receive_message( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: + if not isinstance(message, ValueError): + return + if not stream_error.done(): + stream_error.set_result(message) + # The SDK closes pending requests when its message handler raises. + raise RuntimeError("MCP response stream failed") + # Build session kwargs with optional callbacks session_kwargs: Final[dict[str, Any]] = {} if self._sampling_callback is not None: @@ -456,6 +472,7 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=timedelta(seconds=self.timeout), + message_handler=receive_message if self.transport_type == MCPTransport.http else None, **session_kwargs, ) session: Final = await session_ctx.__aenter__() @@ -467,6 +484,10 @@ class MCPClient: if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) + except McpError: + if stream_error.done(): + raise stream_error.result() + raise finally: try: await session_ctx.__aexit__(None, None, None) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 329dddbdf05..c0ffeaeac62 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -3,12 +3,15 @@ import importlib from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from traceback import walk_tb from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from uuid import uuid4 import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from pydantic import ValidationError from starlette.datastructures import Headers from litellm._logging import verbose_logger @@ -30,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( list_fault_http_status, outcome_wire_value, ) +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree +from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, build_effective_auth_contexts, @@ -78,11 +83,38 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + reference: Final = uuid4().hex + verbose_logger.error( + "MCP connection test failed (reference=%s): %s", + reference, + tuple( + ( + type(cause).__name__, + tuple( + (frame.f_code.co_filename, lineno, frame.f_code.co_name) + for frame, lineno in walk_tb(cause.__traceback__) + ), + ) + for cause in iter_exception_tree(exc) + ), + ) + return next( + ( + message + for cause in iter_exception_tree(exc) + if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None + ), + "An unexpected error occurred while testing the MCP connection. " + f"Retry; if it persists, share reference {reference} with your gateway administrator.", + ) + + +def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None: if isinstance(exc, MCPServerURLCredentialsError): return str(exc.detail) if isinstance(exc, TimeoutError): return ( - f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"Failed to connect to MCP server: no response from {_redact_mcp_resource_url(url) or 'the server'} " f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) @@ -99,10 +131,32 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." - return "Failed to connect to MCP server. Check proxy logs for details." + if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"): + return ( + "Failed to connect to MCP server: the endpoint returned an unsupported content type. " + "Check that the URL is an MCP endpoint, not a web page, and matches the selected transport." + ) + if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"): + return ( + "Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. " + "Check the MCP endpoint URL and the server's protocol implementation." + ) + if MCP_AVAILABLE and isinstance(exc, McpError): + if exc.error.code == 32600 and exc.error.message == "Session terminated": + return ( + "Failed to connect to MCP server: the MCP session was terminated. " + "Check that the URL points to an MCP endpoint and matches the selected transport, " + "then retry to start a new session." + ) + return ( + f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). " + "Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs." + ) + return None if MCP_AVAILABLE: + from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient @@ -1342,7 +1396,6 @@ if MCP_AVAILABLE: except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise except BaseException as e: - verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) return { "status": "error", "error": True, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4db131da62c..49268b4d22d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,9 +1,11 @@ import asyncio import base64 +import json import os import sys from importlib import metadata from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio @@ -11,6 +13,8 @@ import httpx import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from mcp import McpError +from mcp.client.streamable_http import streamable_http_client +from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( LATEST_PROTOCOL_VERSION, @@ -1224,14 +1228,14 @@ def test_without_a_configured_slot_the_existing_precedence_is_unchanged(): _REDIRECT_CASES = [ - ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin + ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin ("https://upstream.example.com/mcp", "https://upstream.example.com:443/other"), # explicit default port - ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host - ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade - ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port - ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host - ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade - ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http + ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host + ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade + ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port + ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host + ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade + ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http ] @@ -1283,3 +1287,109 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: headers = client._get_auth_headers() assert [v for k, v in headers.items() if k.lower() == "esb-oauth"] == ["Bearer minted-token"] assert headers["X-Trace"] == "keep" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("content_type", "body", "expected_type"), + [ + ("text/html", b"secret-page", ValueError), + ("application/json", b"secret-invalid-json", ValidationError), + ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), + ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), + ], +) +async def test_invalid_http_response_surfaces_without_waiting_for_timeout( + content_type: str, body: bytes, expected_type: type[Exception] +) -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": content_type}, content=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(expected_type) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "unsupported content type" in message or "invalid MCP response" in message + assert "secret" not in message + assert "timed out" not in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [200, 401, 503]) +async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": []} + ) + return httpx.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + operation: Final = client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + if status_code == 200: + result: Final = await asyncio.wait_for(operation, timeout=3) + assert result.tools == [] + else: + with pytest.raises(httpx.HTTPStatusError) as caught: + await asyncio.wait_for(operation, timeout=3) + assert caught.value.response.status_code == status_code + + +@pytest.mark.asyncio +async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": "secret-invalid-tools"} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(ValidationError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "invalid MCP response" in message + assert "secret" not in message diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index bd692776c82..73734131130 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3,7 +3,7 @@ import inspect import json import sys from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 @@ -113,7 +113,7 @@ class TestExecuteWithMcpClient: assert "stack_trace" not in result @pytest.mark.asyncio - async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def test_timeout_caps_hanging_operation_and_names_origin(self, monkeypatch): async def fake_create_client(*args, **kwargs): return object() @@ -138,7 +138,7 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] @pytest.mark.asyncio async def test_timeout_covers_client_creation(self, monkeypatch): @@ -166,15 +166,15 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] def test_timeout_defaults_to_tool_listing_timeout(self): default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default assert default == MCP_TOOL_LISTING_TIMEOUT - def test_connection_error_message_timeout_names_url_and_budget(self): + def test_connection_error_message_timeout_names_origin_and_budget(self): message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) - assert "https://api.example.com/mcp/" in message + assert "https://api.example.com" in message assert "30s" in message def test_connection_error_message_hides_arbitrary_http_exception_detail(self): @@ -592,7 +592,7 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert result["error"] is True - assert "Failed to connect to MCP server" in result["message"] + assert "reference" in result["message"] # Error message must not leak raw exception details assert "cancel scope" not in result["message"] @@ -3430,7 +3430,139 @@ class TestConnectionErrorMessage: def test_unknown_error_falls_back_to_generic(self): message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message - assert "proxy logs" in message.lower() + assert "reference" in message.lower() + + def test_sdk_session_terminated_explains_endpoint_and_retry(self) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=32600, message="Session terminated")), "https://example.com/mcp", 30.0 + ) + + assert "session was terminated" in message + assert "MCP endpoint" in message + assert "transport" in message + assert "retry" in message + assert "404" not in message + + @pytest.mark.parametrize("code", [-32700, -32601, -32602, -32603, -32000, 32600, 408]) + def test_rpc_errors_include_code_without_echoing_upstream_data(self, code: int) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=code, message="secret-message", data={"token": "secret-data"})), + "https://example.com/secret-path?token=secret-query", + 30.0, + ) + + assert f"JSON-RPC code {code}" in message + assert "secret" not in message + assert "timed out" not in message + assert "session was terminated" not in message + + @pytest.mark.parametrize("status_code", [401, 403, 404, 405, 429, 503]) + def test_wrapped_http_failures_preserve_status(self, status_code: int) -> None: + response: Final = httpx.Response(status_code, text="secret-body") + upstream: Final = httpx.HTTPStatusError( + "secret-exception", + request=httpx.Request("POST", "https://example.com/?token=secret-query"), + response=response, + ) + wrapped: Final = BaseExceptionGroup( + "secret-group", [asyncio.CancelledError(), BaseExceptionGroup("nested", [upstream])] + ) + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert f"HTTP {status_code}" in message + assert "secret" not in message + + def test_explicit_cause_is_classified_before_incidental_context(self) -> None: + wrapped: Final = RuntimeError("secret-wrapper") + wrapped.__cause__ = httpx.ConnectError("secret-cause") + wrapped.__context__ = TimeoutError("secret-context") + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert "unreachable" in message + assert "secret" not in message + + def test_timeout_url_redacts_credentials_path_query_and_fragment(self) -> None: + message: Final = rest_endpoints._connection_error_message( + TimeoutError("secret-error"), + "https://secret-user:secret-pass@example.com:8443/secret-path?token=secret-query#secret-fragment", + 30.0, + ) + + assert "https://example.com:8443" in message + assert "30s" in message + assert "secret" not in message + + def test_unknown_failure_reference_matches_safe_diagnostics(self, caplog: pytest.LogCaptureFixture) -> None: + import re + + try: + raise RuntimeError("secret-exception-body") + except RuntimeError as exc: + message: Final = rest_endpoints._connection_error_message( + exc, "https://secret-user:secret-password@example.com/secret-path?token=secret-query", 30.0 + ) + + reference: Final = re.search(r"reference ([a-f0-9]{32})", message) + assert reference is not None + diagnostics: Final = tuple( + record for record in caplog.records if "MCP connection test failed" in record.message + ) + assert len(diagnostics) == 1 + assert reference.group(1) in diagnostics[0].message + assert "RuntimeError" in diagnostics[0].message + assert "test_unknown_failure_reference_matches_safe_diagnostics" in diagnostics[0].message + assert diagnostics[0].exc_info is None + assert "secret" not in message + diagnostics[0].message + + @pytest.mark.parametrize("exc", [ValueError("secret-config"), HTTPException(500, "secret-detail")]) + def test_unrelated_errors_are_not_misreported_as_invalid_mcp(self, exc: Exception) -> None: + message: Final = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) + + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + def test_configuration_validation_error_uses_unknown_fallback(self) -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError) as caught: + NewMCPServerRequest.model_validate({"server_name": "example", "transport": "secret-invalid-transport"}) + + message: Final = rest_endpoints._connection_error_message(caught.value, "https://example.com", 30.0) + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + @pytest.mark.asyncio + async def test_connection_test_preserves_cancellation(self) -> None: + async def cancelled_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise asyncio.CancelledError + + payload: Final = NewMCPServerRequest(server_name="cancelled", url="https://example.com", auth_type=MCPAuth.none) + with pytest.raises(asyncio.CancelledError): + await rest_endpoints._execute_with_mcp_client(payload, cancelled_operation) + + @pytest.mark.asyncio + async def test_unknown_failure_preserves_response_contract(self) -> None: + async def failing_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise RuntimeError("secret-operation") + + payload: Final = NewMCPServerRequest(server_name="unknown", url="https://example.com", auth_type=MCPAuth.none) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, failing_operation) + + assert result["error"] is True + assert result["status"] == "error" + assert "reference" in result["message"] + assert "secret" not in result["message"] + assert "stack_trace" not in result class TestGetServerAuthHeaderGroupDefault: From 7b6ef9206e87c09adfc5844c8f8130da73d794fb Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 22:14:47 -0700 Subject: [PATCH 082/157] test(mcp): cover server notifications during tool listing --- .../test_mcp_client.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 49268b4d22d..90292b9b162 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -24,6 +24,7 @@ from mcp.types import ( JSONRPCError, JSONRPCMessage, JSONRPCResponse, + LoggingMessageNotificationParams, ServerCapabilities, ) @@ -1358,6 +1359,58 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co assert caught.value.response.status_code == status_code +@pytest.mark.asyncio +async def test_http_response_handler_preserves_notifications_and_tool_listing() -> None: + notification: Final = { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + } + logging_callback: Final = AsyncMock() + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + if payload["method"] == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"logging": {}, "tools": {}}, + "serverInfo": {"name": "test", "version": "1"}, + }, + }, + ) + response: Final = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}, + } + return httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + content="".join(f"event: message\ndata: {json.dumps(message)}\n\n" for message in (notification, response)), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30, logging_callback=logging_callback) + result: Final = await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ), + timeout=3, + ) + + assert [tool.name for tool in result.tools] == ["search"] + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + + @pytest.mark.asyncio async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message From 884f90c72715b0cfe0d5e2ea09f3744cb8376a78 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 05:15:53 +0000 Subject: [PATCH 083/157] refactor(cost-map): inline background retry, trim tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 119 ++++----- .../test_get_model_cost_map.py | 243 +++++++----------- 2 files changed, 139 insertions(+), 223 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 67d06a0758e..f81ddbfee2e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -184,10 +184,6 @@ def mark_litellm_import_complete() -> None: _litellm_import_complete.set() -def _start_daemon_thread(fn: Callable[[], None]) -> None: - threading.Thread(target=fn, name="litellm-model-cost-map-retry", daemon=True).start() - - @dataclass(frozen=True, slots=True) class ModelCostMapReloaded: model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict @@ -531,32 +527,6 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) -def _use_local_backup(reason: str | None) -> dict: # mutable-ok: returns the mutable model-cost map contract - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = reason - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - - -def _accept_remote( - result: ModelCostMapReloaded, url: str -) -> dict | None: # mutable-ok: returns the mutable model-cost map contract - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, - ) - return None - finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map - _cost_map_source_info.source = "remote" - _cost_map_source_info.url = url - _cost_map_source_info.is_env_forced = False - _cost_map_source_info.fallback_reason = None - return finalized - - def adopt_model_cost_map( new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract ) -> int: @@ -579,7 +549,6 @@ def _retry_remote_fetch_in_background( rng: random.Random, client: _SyncGetClient, first_outcome: _FetchAttemptRetryable, - apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict ) -> None: try: first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) @@ -602,12 +571,20 @@ def _retry_remote_fetch_in_background( ) return _litellm_import_complete.wait() - accepted: Final = _accept_remote(result, url) - if accepted is None: - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) return - apply(accepted) + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + adopt_model_cost_map(finalized) except Exception as e: verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) @@ -619,19 +596,12 @@ def get_model_cost_map( sleep: Callable[[float], None] = time.sleep, rng: random.Random | None = None, client: "_SyncGetClient | None" = None, - start_background: Callable[[Callable[[], None]], None] = _start_daemon_thread, - apply: Callable[ # mutable-ok: injected callback receives the mutable cost-map dict - [dict], - object, - ] = adopt_model_cost_map, ) -> dict: """ Public entry point — returns the model cost map dict. 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. - 2. Otherwise fetches from ``url``, validates the first response, and falls - back to the local backup while retrying transient HTTP errors in the - background. + 2. Otherwise fetches from ``url``, retrying transient errors in a background thread. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -641,9 +611,11 @@ def get_model_cost_map( # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True - return _use_local_backup(None) + _cost_map_source_info.fallback_reason = None + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False @@ -651,32 +623,45 @@ def get_model_cost_map( fetch_client: Final = client if client is not None else httpx fetch_rng: Final = rng if rng is not None else random.Random() outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) - if isinstance(outcome, ModelCostMapReloaded): - accepted: Final = _accept_remote(outcome, url) - return accepted if accepted is not None else _use_local_backup("Remote data failed integrity validation") if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: - verbose_logger.warning( - "LiteLLM: model cost map fetch attempt 1/%d failed: %s; " - "using local backup while retrying in the background", - max_attempts, - outcome.reason, - ) - start_background( - lambda: _retry_remote_fetch_in_background( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=fetch_rng, - client=fetch_client, - first_outcome=outcome, - apply=apply, - ) - ) - else: + threading.Thread( + target=_retry_remote_fetch_in_background, + kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping + "url": url, + "timeout": timeout, + "max_attempts": max_attempts, + "sleep": sleep, + "rng": fetch_rng, + "client": fetch_client, + "first_outcome": outcome, + }, + name="litellm-model-cost-map-retry", + daemon=True, + ).start() + if not isinstance(outcome, ModelCostMapReloaded): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, outcome.reason, ) - return _use_local_backup(f"Remote fetch failed: {outcome.reason}") + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}" + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + content: Final = outcome.model_cost_map + + # Validate using cached count (cheap int comparison, no file I/O) + if not GetModelCostMap.validate_model_cost_map( + fetched_map=content, + backup_model_count=GetModelCostMap._get_backup_model_count(), + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None + return _finalize_loaded_model_cost_map(outcome).model_cost_map diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index e39658a8f37..266c2ca1465 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import threading import pytest @@ -20,7 +21,6 @@ from litellm.litellm_core_utils.get_model_cost_map import ( GetModelCostMap, _count_model_entries, _finalize_model_cost_map, - adopt_model_cost_map, get_model_cost_map_provenance, git_blob_id, ) @@ -566,194 +566,125 @@ from litellm.litellm_core_utils.get_model_cost_map import ( class _SyncSleepRecorder: """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" - def __init__(self): + def __init__(self, block=False): self.waits = [] + self.block = block + self.started = threading.Event() + self.release = threading.Event() def __call__(self, seconds: float) -> None: + if self.block: + self.started.set() + self.release.wait(timeout=10) self.waits.append(seconds) -class _BackgroundRecorder: - def __init__(self): - self.callbacks = [] - - def __call__(self, callback): - self.callbacks.append(callback) +def _retry_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"] -def test_boot_load_returns_local_map_and_schedules_transient_retry(): - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - ], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - - cost_map = get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert len(background.callbacks) == 1 - assert len(cost_map) > 100 - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] is not None - - -def test_background_retry_adopts_valid_remote_map(): - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - httpx.Response(200, content=_real_map_bytes()), - ], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - applied = [] - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - apply=applied.append, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert len(background.callbacks) == 1 - - background.callbacks[0]() - - assert calls["count"] == 2 - assert len(sleeper.waits) == 1 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert len(applied) == 1 - assert applied[0].keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None - - -def test_background_retry_keeps_local_map_after_remaining_failures(): - client, calls = _mock_client( - [httpx.ConnectError("connection refused"), httpx.Response(503)], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - applied = [] - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - apply=applied.append, - ) - background.callbacks[0]() - - assert calls["count"] == 3 - assert len(sleeper.waits) == 2 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert 4.0 <= sleeper.waits[1] < 5.0 - assert applied == [] - assert get_model_cost_map_source_info()["source"] == "local" - - -def test_boot_load_does_not_schedule_non_retryable_failure(): - client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - ) - - assert calls["count"] == 1 - assert sleeper.waits == [] - assert background.callbacks == [] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] is not None - - -def test_boot_load_success_does_not_schedule_background_retry(): +def test_boot_load_success_does_not_start_background_retry(): client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() cost_map = get_model_cost_map( url=_URL, sleep=sleeper, rng=random.Random(0), client=client, - start_background=background, ) assert calls["count"] == 1 assert sleeper.waits == [] - assert background.callbacks == [] - assert get_model_cost_map_source_info()["source"] == "remote" + assert _retry_threads() == [] assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + assert get_model_cost_map_source_info()["source"] == "remote" -def test_boot_load_with_one_attempt_does_not_schedule_background_retry(): - client, calls = _mock_client([httpx.ConnectError("connection refused")], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): + import litellm + from litellm import utils as litellm_utils + from litellm.litellm_core_utils import get_model_cost_map as module - get_model_cost_map( + original_model_cost = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + for name, provider_models in tuple(vars(litellm).items()): + if name.endswith("_models") and isinstance(provider_models, set): + monkeypatch.setattr(litellm, name, set(provider_models)) + monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + source_info = module._cost_map_source_info + for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): + monkeypatch.setattr(source_info, name, getattr(source_info, name)) + + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} + client, calls = _mock_client( + [ + httpx.ConnectError("connection refused"), + httpx.Response(200, content=json.dumps(remote_map).encode()), + ], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder(block=True) + litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) + + cost_map = get_model_cost_map( url=_URL, - max_attempts=1, + max_attempts=3, sleep=sleeper, rng=random.Random(0), client=client, - start_background=background, ) assert calls["count"] == 1 assert sleeper.waits == [] - assert background.callbacks == [] - assert get_model_cost_map_source_info()["source"] == "local" - - -def test_adopt_model_cost_map_replays_runtime_registration_and_provider_models(): - import litellm - from litellm import utils as litellm_utils - - original_model_cost = litellm.model_cost - original_registry = dict(litellm_utils._runtime_registered_model_cost) - original_anthropic_models = set(litellm.anthropic_models) + assert sleeper.started.wait(timeout=10) + threads = _retry_threads() try: - litellm.register_model( - model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} - ) - - models_count = adopt_model_cost_map({"anthropic/new-model": {"litellm_provider": "anthropic", "mode": "chat"}}) - - assert models_count == 1 - assert "anthropic/new-model" in litellm.anthropic_models - assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 + assert len(threads) == 1 + assert "claude-remote-only-test" not in cost_map + assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed:") + sleeper.release.set() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0 + assert calls["count"] == 2 + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + assert "my-runtime-model" in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None finally: - litellm.model_cost = original_model_cost # test-quality-ok: restore the module state changed by adoption - litellm_utils._runtime_registered_model_cost.clear() - litellm_utils._runtime_registered_model_cost.update(original_registry) - litellm.anthropic_models.clear() - litellm.anthropic_models.update(original_anthropic_models) - litellm_utils._invalidate_model_cost_lowercase_map() + sleeper.release.set() + for thread in _retry_threads(): + thread.join(timeout=10) + + +def test_boot_load_does_not_retry_non_retryable_failure(): + client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert _retry_threads() == [] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None def test_boot_load_respects_local_env_override(monkeypatch): From 0721163cacfbd9fbfeee5dd25205cb76840111dd Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:49:40 -0700 Subject: [PATCH 084/157] test(e2e/ui): cover team-scoped model visibility, re-editing litellm params, and model health checks (#40039) * test(e2e/ui): cover team-scoped model visibility, re-editing litellm params, and model health checks Three Models and Endpoints flows had no end-to-end coverage, and all three keep coming back as bug reports. modelsByTeam walks an internal user through the Current team control and asserts the table lists exactly what each team grants. It creates one deployment that belongs to no team, proves that deployment is visible under Personal, then proves it is absent under both seeded teams, so an empty table cannot pass the same assertions. editLitellmParams adds a temperature and a custom pair to a deployment, saves, then re-edits the temperature and drops the custom pair. It checks both update request bodies, polls the stored deployment until the new temperature is there, reloads the page to confirm the second save is what renders, and sends one chat completion to prove the deployment still serves. modelHealthStatus runs the health check on a reachable deployment and on one pointed at a dead port, asserts the healthy and unhealthy cells and the two detail dialogs, and reloads to confirm both statuses are stored. Every deployment these specs create carries a unique name and is deleted in afterEach, including on the failure path. * test(e2e/ui): find health rows across every page of the health table The health table pages server-side at 50 rows with no search box, so on a proxy carrying more deployments than that the two deployments the spec creates can land on a later page and the lookup finds nothing. Row lookups now walk the pages, using the table's own page indicator to know when to advance and when to wrap back to the first page. * test(e2e/ui): build the created deployment ids without mutating the array * test(ui): scope model deployments to Playwright fixtures --- .../tests/internal-user/modelsByTeam.spec.ts | 196 ++++++++++++++ .../modelsPage/editLitellmParams.spec.ts | 252 ++++++++++++++++++ .../modelsPage/modelHealthStatus.spec.ts | 245 +++++++++++++++++ 3 files changed, 693 insertions(+) create mode 100644 tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts create mode 100644 tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts create mode 100644 tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts diff --git a/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts new file mode 100644 index 00000000000..5e2c80b5845 --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts @@ -0,0 +1,196 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_ORG_ALIAS, + INTERNAL_USER_STORAGE_PATH, +} from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, CHAT_MODEL_B, masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CURRENT_TEAM_VIEW = "Current Team Models"; +const ALL_MODELS_VIEW = "All Available Models"; +const PERSONAL_TEAM = "Personal"; + +const teamSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "Current team", exact: true }); +const viewSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "View", exact: true }); + +async function chooseOption( + page: PlaywrightPage, + selector: Locator, + optionName: string, +): Promise { + await selector.click(); + const option = page.getByRole("option", { name: optionName, exact: true }); + await expect(option, `option ${optionName} is offered`).toBeVisible({ + timeout: 10_000, + }); + await option.click(); + await expect( + selector, + `${optionName} is the selection the control now reports`, + ).toContainText(optionName, { + timeout: 10_000, + }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +function modelRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ ungrantedModelName: string }>({ + ungrantedModelName: async ({ page }, use) => { + const ungrantedModelName = `e2e-ungranted-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: ungrantedModelName, + litellm_params: { + model: `openai/${ungrantedModelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const ungrantedModelId = (await created.json()).model_info?.id; + expect(ungrantedModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll(async () => await isRegistered(page, ungrantedModelName), { + message: `deployment ${ungrantedModelName} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await use(ungrantedModelName); + } finally { + await deleteDeployment(page, ungrantedModelId); + } + }, +}); + +test.describe("Models and Endpoints for an internal user", () => { + test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); + + test("shows an internal user exactly the models of the team they select", async ({ + page, + ungrantedModelName, + }) => { + await navigateToPage(page, Page.Models); + + await expect( + page.getByRole("tab", { name: "Your Models" }), + "an internal user lands on their own models tab, not an admin-only view", + ).toBeVisible({ timeout: 15_000 }); + await expect( + viewSelector(page), + "the models table opens scoped to the selected team", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + `the personal view lists ${ungrantedModelName}, so it is on the proxy and reachable from this page`, + ).toHaveCount(1, { timeout: 30_000 }); + + await chooseOption(page, teamSelector(page), E2E_TEAM_CRUD_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_B}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is on the proxy but not granted to ${E2E_TEAM_CRUD_ALIAS}, so it must not be listed`, + ).toHaveCount(0); + + await chooseOption(page, teamSelector(page), E2E_TEAM_ORG_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_ORG_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + page.getByTestId("pagination-range"), + `${E2E_TEAM_ORG_ALIAS} lists the one model it grants and nothing else`, + ).toHaveText("Showing 1-1 of 1", { timeout: 15_000 }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${CHAT_MODEL_B} belongs to another team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is granted to no team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + + await chooseOption(page, viewSelector(page), ALL_MODELS_VIEW); + await expect( + modelRow(page, CHAT_MODEL_A), + `switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`, + ).toHaveCount(1, { timeout: 15_000 }); + + await page.reload(); + await expect( + teamSelector(page), + "the team selection is not persisted across a reload, so the table returns to the personal view", + ).toContainText(PERSONAL_TEAM, { timeout: 15_000 }); + await expect( + viewSelector(page), + "the view selection is not persisted across a reload either", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + "the personal view still renders models after a reload rather than coming back empty", + ).toHaveCount(1, { timeout: 30_000 }); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts new file mode 100644 index 00000000000..4480515ae59 --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts @@ -0,0 +1,252 @@ +import { + test as base, + expect, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { masterKey, sendChatCompletion } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CUSTOM_PARAM = "extra_headers"; +const CUSTOM_PARAM_VALUE = { "X-E2E-Edit-Probe": "one" }; + +type StoredParams = Record; + +async function readStoredParams( + page: PlaywrightPage, + modelId: string, +): Promise { + const body = await readBack<{ data: { litellm_params: StoredParams }[] }>( + page, + `/model/info?litellm_model_id=${modelId}`, + ); + return body.data[0]?.litellm_params ?? {}; +} + +function paramsEditor(page: PlaywrightPage) { + return page.getByPlaceholder('"rpm": 100'); +} + +async function editParams( + page: PlaywrightPage, + mutate: (params: StoredParams) => StoredParams, +): Promise { + await page.getByRole("button", { name: "Edit Settings" }).click(); + const editor = paramsEditor(page); + await expect( + editor, + "the LiteLLM Params editor is reachable on every visit to the edit form", + ).toBeVisible({ + timeout: 15_000, + }); + const shown = JSON.parse(await editor.inputValue()) as StoredParams; + await editor.fill(JSON.stringify(mutate(shown), null, 2)); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ + deployment: { readonly modelName: string; readonly createdModelId: string }; +}>({ + deployment: async ({ page, request }, use) => { + const modelName = `e2e-edit-params-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: modelName, + litellm_params: { + model: `openai/${modelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const createdModelId = (await created.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll( + async () => { + try { + await sendChatCompletion(request, { + model: modelName, + prompt: `warmup ${modelName}`, + }); + return true; + } catch { + return false; + } + }, + { + message: `deployment ${modelName} never became routable after /model/new`, + timeout: 60_000, + }, + ) + .toBe(true); + await use({ modelName, createdModelId }); + } finally { + await deleteDeployment(page, createdModelId); + } + }, +}); + +test.describe("Edit LiteLLM Params on a deployment", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("params added on a deployment can be re-edited, and the deployment keeps serving", async ({ + page, + request, + deployment: { modelName, createdModelId }, + }) => { + await navigateToPage(page, Page.Models); + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect( + modelIdCell, + `the Models table lists ${modelName}`, + ).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ + timeout: 15_000, + }); + + await editParams(page, (params) => ({ + ...params, + temperature: 0.2, + [CUSTOM_PARAM]: CUSTOM_PARAM_VALUE, + })); + const firstSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + firstSave.litellm_params?.temperature, + "the added temperature goes on the wire", + ).toBe(0.2); + expect( + firstSave.litellm_params?.[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} goes on the wire`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + firstSave.litellm_params?.model, + "a params edit does not rewrite the upstream model", + ).toBe(`openai/${modelName}`); + expect( + firstSave.litellm_params?.api_base, + "a params edit does not rewrite the api base", + ).toBe(MOCK_LLM_BASE); + expect( + firstSave.litellm_params, + "the credential is never re-sent, so a masked placeholder cannot overwrite the stored key", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: "the added temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.2); + const afterFirstSave = await readStoredParams(page, createdModelId); + expect( + afterFirstSave[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} reached the stored deployment`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + afterFirstSave.model, + "the stored upstream model survived the edit", + ).toBe(`openai/${modelName}`); + expect( + afterFirstSave.api_base, + "the stored api base survived the edit", + ).toBe(MOCK_LLM_BASE); + + await editParams(page, (params) => ({ + ...Object.fromEntries( + Object.entries(params).filter(([key]) => key !== CUSTOM_PARAM), + ), + temperature: 0.7, + })); + const secondSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + secondSave.litellm_params?.temperature, + "a param set by an earlier save can be edited again", + ).toBe(0.7); + expect( + secondSave.litellm_params, + `dropping ${CUSTOM_PARAM} from the editor drops it from the request the UI sends`, + ).not.toHaveProperty(CUSTOM_PARAM); + expect( + secondSave.litellm_params?.model, + "a second params edit still leaves the upstream model alone", + ).toBe(`openai/${modelName}`); + expect( + secondSave.litellm_params?.api_base, + "a second params edit still leaves the api base alone", + ).toBe(MOCK_LLM_BASE); + expect( + secondSave.litellm_params, + "the credential is still never re-sent", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: + "the re-edited temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.7); + + await page.reload(); + await expect( + page + .getByRole("tabpanel", { name: "Overview" }) + .getByText('"temperature": 0.7'), + "reopening the deployment renders the re-edited value, not the one from the first save", + ).toBeVisible({ timeout: 20_000 }); + + await sendChatCompletion(request, { + model: modelName, + prompt: `still serving ${modelName}`, + }); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts new file mode 100644 index 00000000000..247cce1b85d --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts @@ -0,0 +1,245 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const UNREACHABLE_BASE = "http://127.0.0.1:9/v1"; + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +function healthRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +function pageOf(label: string): { current: number; total: number } { + const [current, total] = label + .replace("Page ", "") + .split(" of ") + .map((part) => Number(part.trim())); + return { current, total }; +} + +async function locateHealthRow( + page: PlaywrightPage, + modelName: string, +): Promise { + const pageLabel = page.getByTestId("pagination-page"); + await expect( + pageLabel, + "the health table reports which page it is showing", + ).toBeVisible({ timeout: 20_000 }); + + const deadline = Date.now() + 60_000; + while (Date.now() < deadline) { + const row = healthRow(page, modelName); + const onThisPage = await row + .first() + .waitFor({ state: "visible", timeout: 3_000 }) + .then(() => true) + .catch(() => false); + if (onThisPage) return row; + + const { current, total } = pageOf(await pageLabel.innerText()); + const goTo = current < total ? current + 1 : 1; + if (total === 1) continue; + await page + .getByRole("button", { + name: current < total ? "Go to next page" : "Go to first page", + }) + .click(); + await expect(pageLabel).toContainText(`Page ${goTo} of`, { + timeout: 15_000, + }); + } + return healthRow(page, modelName); +} + +async function openHealthTab(page: PlaywrightPage): Promise { + await page.getByRole("tab", { name: "Health Status" }).click(); + await expect( + page.getByRole("heading", { name: "Model Health Status" }), + ).toBeVisible({ timeout: 15_000 }); +} + +async function expectStatus( + page: PlaywrightPage, + modelName: string, + status: string, +): Promise { + const row = await locateHealthRow(page, modelName); + await expect(row, `${modelName} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await expect( + row.getByText(status, { exact: true }), + `the Health Status cell for ${modelName} reads ${status}`, + ).toHaveCount(1, { timeout: 60_000 }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +async function withDeployment( + page: PlaywrightPage, + prefix: string, + apiBase: string, + use: (name: string) => Promise, +): Promise { + const name = `${prefix}-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: name, + litellm_params: { + model: `openai/${name}`, + api_base: apiBase, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new for ${name} failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const id = (await created.json()).model_info?.id; + expect(id, `model id from /model/new for ${name}`).toBeTruthy(); + try { + await expect + .poll(() => isRegistered(page, name), { + message: `deployment ${name} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await use(name); + } finally { + await deleteDeployment(page, id); + } +} + +const test = base.extend<{ reachableName: string; unreachableName: string }>({ + reachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-up", MOCK_LLM_BASE, use); + }, + unreachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-down", UNREACHABLE_BASE, use); + }, +}); + +test.describe("Model health status", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Run Health Check reports a reachable deployment healthy and an unreachable one unhealthy", async ({ + page, + reachableName, + unreachableName, + }) => { + await navigateToPage(page, Page.Models); + await openHealthTab(page); + + for (const name of [reachableName, unreachableName]) { + const row = await locateHealthRow(page, name); + await expect(row, `${name} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await row + .getByRole("button", { name: "Run Health Check", exact: true }) + .click(); + } + + await expectStatus(page, reachableName, "healthy"); + await expect( + healthRow(page, reachableName).getByText("unhealthy", { exact: true }), + "a reachable deployment is never reported unhealthy", + ).toHaveCount(0); + await expectStatus(page, unreachableName, "unhealthy"); + + const successDetail = ( + await locateHealthRow(page, reachableName) + ).getByRole("button", { + name: "View response details", + }); + await expect( + successDetail, + `${reachableName} offers its health check response for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await successDetail.click(); + const successDialog = page.getByRole("dialog"); + await expect( + successDialog.getByRole("heading", { + name: `Health Check Response - ${reachableName}`, + }), + "the healthy deployment's detail opens its own response dialog", + ).toBeVisible({ timeout: 10_000 }); + await successDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(successDialog).toBeHidden({ timeout: 10_000 }); + + const errorDetail = ( + await locateHealthRow(page, unreachableName) + ).getByRole("button", { + name: "View full error details", + }); + await expect( + errorDetail, + `${unreachableName} offers its health check error for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await errorDetail.click(); + const errorDialog = page.getByRole("dialog"); + await expect( + errorDialog.getByRole("heading", { + name: `Health Check Error - ${unreachableName}`, + }), + "the unreachable deployment's detail opens its own error dialog", + ).toBeVisible({ timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog carries the upstream connection failure, not a generic message", + ).toContainText(/connection error/i, { timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog names the endpoint that could not be reached", + ).toContainText(UNREACHABLE_BASE); + await errorDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(errorDialog).toBeHidden({ timeout: 10_000 }); + + await page.reload(); + await openHealthTab(page); + await expectStatus(page, reachableName, "healthy"); + await expectStatus(page, unreachableName, "unhealthy"); + }); +}); From b3151073d2a8ee274fd067c1d4cfb17da0a82459 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:49:54 -0700 Subject: [PATCH 085/157] test(e2e/ui): cover key budget window, non-admin model scope edit, and key blocking (#40027) * test(e2e/ui): cover key budget window, non-admin model scope edit, and key blocking Three Playwright specs for the Virtual Keys flows customers hit most, each reading its result back through /key/info and /v1/chat/completions rather than trusting the toast: - a monthly spend cap and reset window set through Edit Settings, surviving a reload, with clearing the window leaving the cap in place - a team member narrowing their own team key's models, and the proxy refusing the model they dropped - blocking a key from its detail page, then unblocking it Each test owns the key it edits and deletes it on teardown, so retries and --repeat-each never run out of fixtures. * test(e2e/ui): tighten virtual key specs from review feedback Replace the mutable suite-level key state with a Playwright fixture, so the alias and token are never reassigned and cleanup stays tied to the test. Assert /key/delete succeeded instead of discarding the response, so a failed cleanup surfaces rather than leaving rows behind. Drop the explanatory JSDoc the repo's comment policy disallows, keeping only the one line explaining why Date.now() alone is not unique enough. Type the master-key POST helper against a real guard instead of casting to Record. Assert the unblocked key is served with a 200, not just the response text, and that clearing the reset window also clears budget_reset_at. * test(ui): assert the team response through Playwright --- tests/e2e/ui/helpers/navigation.ts | 10 + tests/e2e/ui/helpers/traffic.ts | 51 ++++- .../internalUserKeyScope.spec.ts | 208 ++++++++++++++++++ .../ui/tests/proxy-admin/keyBlocking.spec.ts | 112 ++++++++++ .../tests/proxy-admin/keyBudgetWindow.spec.ts | 101 +++++++++ 5 files changed, 478 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 4a7c4e7baa9..e0e7b4da396 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -73,3 +73,13 @@ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise await cell.click(); await expect(page.getByText("Back to Teams")).toBeVisible({ timeout: 10_000 }); } + +export async function openKeyDetail(page: PlaywrightPage, alias: string): Promise { + await page.getByPlaceholder("Search by key alias or ID").fill(alias); + const row = page.getByRole("row").filter({ hasText: alias }); + await expect(row, `key row "${alias}" never appeared on the Virtual Keys page`).toBeVisible({ timeout: 15_000 }); + await row.getByRole("button", { name: alias }).click(); + await expect(page.getByText("Back to Keys"), `key detail for "${alias}" never opened`).toBeVisible({ + timeout: 15_000, + }); +} diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 7f8417cdffb..cb68747b364 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -1,4 +1,4 @@ -import { APIRequestContext, expect } from "@playwright/test"; +import { APIRequestContext, APIResponse, expect } from "@playwright/test"; /** Model names served by fixtures/config.yml, both backed by the mock LLM server. */ export const CHAT_MODEL_A = "fake-openai-gpt-4"; @@ -15,6 +15,9 @@ export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-123 export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; +/** Date.now() alone collides: `--repeat-each` starts its copies inside the same millisecond. */ +export const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + interface ChatOptions { model: string; prompt: string; @@ -25,9 +28,8 @@ interface ChatOptions { traceId?: string; } -/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ -export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { - const res = await request.post(`${rootPath()}/v1/chat/completions`, { +const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise => + request.post(`${rootPath()}/v1/chat/completions`, { headers: { Authorization: `Bearer ${opts.apiKey ?? masterKey()}`, "Content-Type": "application/json", @@ -39,12 +41,26 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}), }, }); + +/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ +export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); const body = await res.json(); expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); return body.id as string; } +export interface ChatAttempt { + status: number; + body: string; +} + +export async function attemptChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); + return { status: res.status(), body: await res.text() }; +} + /** `key` is the sk- value to authenticate with; `token` is its hash, which spend aggregates are keyed by. */ export async function createVirtualKey( request: APIRequestContext, @@ -66,6 +82,33 @@ export async function createVirtualKey( }; } +export interface KeyInfo { + key_alias: string | null; + max_budget: number | null; + budget_duration: string | null; + budget_reset_at: string | null; + blocked: boolean | null; + models: string[]; + team_id: string | null; +} + +export async function readKeyInfo(request: APIRequestContext, token: string): Promise { + const res = await request.get(`${rootPath()}/key/info?key=${encodeURIComponent(token)}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /key/info for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); + const body = await res.json(); + return body.info as KeyInfo; +} + +export async function deleteVirtualKey(request: APIRequestContext, token: string): Promise { + const res = await request.post(`${rootPath()}/key/delete`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { keys: [token] }, + }); + expect(res.ok(), `key delete for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); +} + /** Spend logs are flushed on a timer, so an assertion straight after a completion races the writer. */ export async function waitForSpendLog( request: APIRequestContext, diff --git a/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts new file mode 100644 index 00000000000..f923841257a --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts @@ -0,0 +1,208 @@ +import { test, expect, type APIRequestContext } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { + dismissFeedbackPopup, + navigateToPage, + openKeyDetail, +} from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + masterKey, + readKeyInfo, + rootPath, + uniqueSuffix, +} from "../../helpers/traffic"; + +const MEMBER_PASSWORD = "E2e-Team-Member-Pass-1!"; + +interface CreatedTeam { + readonly team_id: string; +} + +function assertCreatedTeam(body: unknown): asserts body is CreatedTeam { + expect(body, "/team/new returned no team_id").toMatchObject({ + team_id: expect.any(String), + }); +} + +async function postAsMaster( + request: APIRequestContext, + path: string, + data: Record, +): Promise { + const res = await request.post(`${rootPath()}${path}`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data, + }); + expect( + res.ok(), + `POST ${path} failed (${res.status()}): ${await res.text()}`, + ).toBe(true); + return res.json(); +} + +test.describe("Internal User - own team key model scope", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("a team member narrows their own key's models and the proxy enforces it", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const email = `team-member-${suffix}@test.local`; + const userId = `e2e-key-scope-user-${suffix}`; + const alias = `e2e-key-scope-${suffix}`; + + const team = await postAsMaster(request, "/team/new", { + team_alias: `E2E Key Scope ${suffix}`, + models: [CHAT_MODEL_A, CHAT_MODEL_B], + team_member_permissions: ["/key/generate", "/key/update", "/key/info"], + }); + assertCreatedTeam(team); + const teamId = team.team_id; + + try { + await postAsMaster(request, "/user/new", { + user_id: userId, + user_email: email, + user_role: "internal_user", + auto_create_key: false, + }); + await postAsMaster(request, "/user/update", { + user_id: userId, + password: MEMBER_PASSWORD, + }); + await postAsMaster(request, "/team/member_add", { + team_id: teamId, + member: { role: "user", user_id: userId }, + }); + + const created = await createVirtualKey(request, { + key_alias: alias, + team_id: teamId, + user_id: userId, + models: [], + }); + + try { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page + .getByPlaceholder("Enter your password") + .fill(MEMBER_PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect( + page.locator("a", { hasText: "Virtual Keys" }), + `${email} never reached the dashboard`, + ).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await expect( + page.getByRole("option", { name: CHAT_MODEL_A, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_A} to a team member`, + ).toBeVisible({ timeout: 15_000 }); + await expect( + page.getByRole("option", { name: CHAT_MODEL_B, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_B} to a team member`, + ).toBeVisible(); + + await page + .getByRole("option", { name: CHAT_MODEL_A, exact: true }) + .click(); + await page.keyboard.press("Escape"); + + const updated = page.waitForResponse( + (res) => + res.url().includes("/key/update") && + res.request().method() === "POST", + ); + await page.getByRole("button", { name: "Save Changes" }).click(); + const updateStatus = (await updated).status(); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeGreaterThanOrEqual(200); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeLessThan(300); + await expect( + page.getByText("Key updated successfully").first(), + ).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => (await readKeyInfo(request, created.token)).models, + { + message: `the narrowed model scope never reached /key/info for ${alias}`, + timeout: 20_000, + }, + ) + .toEqual([CHAT_MODEL_A]); + + await expect + .poll( + async () => + await attemptChatCompletion(request, { + model: CHAT_MODEL_B, + prompt: `out of scope ${suffix}`, + apiKey: created.key, + }), + { + message: `${CHAT_MODEL_B} was still served after the key was narrowed to ${CHAT_MODEL_A}`, + timeout: 30_000, + }, + ) + .toMatchObject({ + status: 403, + body: expect.stringContaining(CHAT_MODEL_B), + }); + + const inScope = await attemptChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `in scope ${suffix}`, + apiKey: created.key, + }); + expect( + inScope, + `${CHAT_MODEL_A} is no longer served by the narrowed key`, + ).toMatchObject({ + status: 200, + body: expect.stringContaining(MOCK_RESPONSE_TEXT), + }); + } finally { + await deleteVirtualKey(request, created.token); + } + } finally { + await request.post(`${rootPath()}/user/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { user_ids: [userId] }, + }); + await request.post(`${rootPath()}/team/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { team_ids: [teamId] }, + }); + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts new file mode 100644 index 00000000000..99a8065a797 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts @@ -0,0 +1,112 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + readKeyInfo, + sendChatCompletion, + uniqueSuffix, +} from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; + apiKey: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-block-key-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token, apiKey: created.key }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key blocking", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("blocking a key stops it serving and unblocking restores it", async ({ page, scopedKey }) => { + const { alias, token, apiKey } = scopedKey; + + await sendChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: `pre-block ${alias}`, + apiKey, + }); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Block Key" }).click(); + const blockDialog = page.getByRole("dialog", { name: "Block Key" }); + await expect(blockDialog, "the Block Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await blockDialog.getByRole("button", { name: "Block", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back blocked from /key/info", + timeout: 20_000, + }) + .toBe(true); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "blocked", + apiKey, + }), + { + message: "a blocked key was still served by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 401, body: expect.stringContaining("blocked") }); + + await page.reload(); + await expect( + page.getByText("Blocked", { exact: true }), + "the reloaded key detail does not show the key as blocked", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Unblock Key" }).click(); + const unblockDialog = page.getByRole("dialog", { name: "Unblock Key" }); + await expect(unblockDialog, "the Unblock Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await unblockDialog.getByRole("button", { name: "Unblock", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back unblocked from /key/info", + timeout: 20_000, + }) + .toBe(false); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "unblocked", + apiKey, + }), + { + message: "an unblocked key is still refused by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 200, body: expect.stringContaining(MOCK_RESPONSE_TEXT) }); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts new file mode 100644 index 00000000000..4e4d0a395c3 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -0,0 +1,101 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, createVirtualKey, deleteVirtualKey, readKeyInfo, uniqueSuffix } from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-budget-window-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + team_id: E2E_TEAM_CRUD_ID, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key budget window", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => { + const { alias, token } = scopedKey; + + const before = await readKeyInfo(page.request, token); + expect(before.max_budget, "a freshly generated key starts with no budget").toBeNull(); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "monthly", exact: true }).click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { + message: "the $12.50 cap never reached /key/info", + timeout: 20_000, + }) + .toBe(12.5); + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the monthly reset window never reached /key/info", + timeout: 20_000, + }) + .toBe("30d"); + + const capped = await readKeyInfo(page.request, token); + const resetAt = new Date(capped.budget_reset_at ?? ""); + expect(Number.isNaN(resetAt.getTime()), "a monthly window left the key with no budget_reset_at").toBe(false); + expect(resetAt.getTime(), "budget_reset_at was set in the past").toBeGreaterThan(Date.now()); + expect(resetAt.getUTCDate(), "a monthly window resets on the 1st, a daily one would not").toBe(1); + + await page.reload(); + await expect( + page.getByRole("paragraph").filter({ hasText: "of $12.50" }), + "the reloaded key detail does not render the $12.50 cap", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("tab", { name: "Settings" }).click(); + await expect( + page.getByTestId("budget-reset-value"), + "the reloaded key detail does not name the 30d reset window", + ).toHaveText(/Every 30d/, { timeout: 15_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "Never resets", exact: true }).click(); + + const cleared = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(cleared).toHaveProperty("budget_duration"); + expect(cleared.budget_duration, "clearing the window must send budget_duration: null explicitly").toBeNull(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the reset window was never cleared on /key/info", + timeout: 20_000, + }) + .toBeNull(); + + const after = await readKeyInfo(page.request, token); + expect(after.budget_reset_at, "clearing the reset window left a stale next-reset timestamp").toBeNull(); + expect(after.max_budget, "clearing the reset window also wiped the spend cap").toBe(12.5); + expect(after.models, "editing the budget left the key's models untouched").toEqual(before.models); + expect(after.team_id, "editing the budget left the key's team untouched").toEqual(before.team_id); + }); +}); From 36bd7f113837caadd3f3d400dd246cd4649b7b33 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:50:13 -0700 Subject: [PATCH 086/157] fix(mcp): honor an explicit null on toolset update, cover MCP lifecycle e2e (#40022) * fix(mcp): honor an explicit null on toolset update, cover MCP lifecycle e2e PUT /v1/mcp/toolset dumped its payload with exclude_none, so a field sent as null looked exactly like one the caller left out and the stored value survived. An admin could not clear a toolset's description: the save reported success and the old text came straight back. It now dumps with exclude_unset, so absent keeps and null clears, which is what PUT /v1/mcp/server already did. A null tools list clears the selection to empty, and a null toolset_name is ignored because a toolset always has a name. Adds create, read, partial-update, clear and delete e2e coverage for MCP servers and toolsets, with every read-back polled on every replica so an edit that lands on one replica and not another fails the test, plus an enforcement test proving a key granted a toolset lists exactly that toolset's tools against the real Datadog upstream. * fix(e2e): refuse a read-back that no replica serves A read-back over an empty replica mapping satisfied every predicate and returned as if it had converged, so it would have asserted nothing and passed. No wiring can produce that today, since the replica list always falls back to at least one URL, but a helper whose whole job is proving a write reached every replica should not have a shape that passes vacuously. * fix(mcp): keep a null tools list a no-op on toolset update Treating a null tools list as a clear meant an existing client that sends tools=null during a partial update, meaning "leave the selection alone", silently lost every tool the toolset grants. That is a permission surface, so the quiet version of it is the worst version. A toolset always has a tool list, the same way it always has a name, so a null on either is now a no-op. Emptying the selection is an explicit [], which cannot be confused with a field the caller left out, and which is what the dashboard already sends. * fix(e2e): keep MCP admin routes on the data plane /v1/mcp/* is a lazily mounted feature, so a gateway registers it on the first matching request, which happens after the startup route trim that drops management endpoints. Routing it to the control plane therefore sent every MCP call to the one backend process: the new lifecycle read-backs proved a single process rather than every replica, and mcp_client's await_registered barrier waited on a registry that does not serve the tools/list call it guards, so the existing MCP suites polled a gateway that had not synced yet until poll_timeout Verified against a two-gateway split stack (backend on 4001, gateways on 4010 and 4011, one postgres): both gateways answer /v1/mcp/server and /v1/mcp/toolset, and each served 6 server reads and 7 toolset reads over the run * fix(e2e): grant the toolset by the tool's own name, not the wire name tools/list serves a tool as , but a toolset grants by the tool's own name: resolve_toolset_permissions reads toolset.tools[].tool_name straight through, and the prefix is added on the way out. The test built the toolset from the names tools/list reported, so the grant matched nothing, the scoped key listed no tools, and await_tools ran out its whole poll_timeout before failing Measure the prefix off search_datadog_logs, whose own name is known, rather than guessing it from the alias, since the proxy can be configured to prefix with a short server id instead. The expectation compared against tools/list stays in wire names; only what the toolset stores crosses back * test(mcp): build immutable lifecycle updates and replica results * test: validate opaque stream IDs and hide log-reader credentials * test: isolate auto-router scenarios and clean partial setup * test: honor Datadog search rate-limit reset headers * test: share the Datadog read-back deadline across retries * test: preserve captured MCP toolset update fields --- .../_experimental/mcp_server/toolset_db.py | 16 +- .../mcp_management_endpoints.py | 4 + tests/e2e/coverage_registry/mcp.yaml | 8 + tests/e2e/coverage_registry/mgmt.yaml | 10 + tests/e2e/e2e_http.py | 64 +++- .../test_responses_bridge_streaming_e2e.py | 24 +- tests/e2e/logging/datadog_reader.py | 124 +++++--- tests/e2e/logging/test_datadog_reader.py | 223 +++++++++++++ tests/e2e/management/management_client.py | 35 +++ .../e2e/management/test_mcp_lifecycle_e2e.py | 294 ++++++++++++++++++ tests/e2e/mcp/datadog_mcp.py | 7 +- tests/e2e/mcp/mcp_client.py | 39 ++- .../mcp/test_mcp_toolset_enforcement_e2e.py | 95 ++++++ tests/e2e/models.py | 73 ++++- tests/e2e/proxy_client.py | 216 ++++++++++++- .../test_auto_router_regressions_e2e.py | 239 +++++++------- tests/e2e/test_e2e_http.py | 91 +++++- tests/e2e/test_proxy_client.py | 95 +++++- .../mcp_server/test_mcp_partial_update.py | 71 ++++- 19 files changed, 1501 insertions(+), 227 deletions(-) create mode 100644 tests/e2e/logging/test_datadog_reader.py create mode 100644 tests/e2e/management/test_mcp_lifecycle_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index ecaaf35e817..48bad178927 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -132,10 +132,18 @@ async def update_mcp_toolset( data: UpdateMCPToolsetRequest, touched_by: str, ) -> MCPToolset | None: - data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"}) - if "tools" in data_dict: - data_dict["tools"] = json.dumps(data_dict["tools"]) - data_dict["updated_by"] = touched_by + """A partial update: absent keeps, null clears. A toolset always has a name and a + tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear; + emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a + caller that left the field out.""" + data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization + ( + (field, json.dumps(value) if field == "tools" else value) + for field, value in data.model_dump(exclude_unset=True).items() + if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) + ), + updated_by=touched_by, + ) try: row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d5c3427f29a..2aa7fdc7393 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2673,6 +2673,8 @@ if MCP_AVAILABLE: """ Updates the MCP Server in the db. + Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared. + Parameters: - payload: UpdateMCPServerRequest - Required. The updated mcp server data. ``` @@ -3098,6 +3100,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: str | None = Header(None), ): + """Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except + ``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit [].""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..f853e9ff8d6 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,11 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "user_api_key_auth_mcp.py:2137" + rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 860d96a50b4..c8d7037d2fd 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,3 +76,13 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} +- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} +- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} +- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"} +- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"} +- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"} +- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"} +- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} +- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} +- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..415c72bbb3c 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders): anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version") +class PartialBody(BaseModel): + """A body for a partial-update route (absent = keep, null = clear): a field left + unset is omitted from the wire, and a field set to None is sent as JSON null.""" + + class NoBody(BaseModel): """Empty body/query for routes that take none.""" @@ -252,6 +257,13 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) + +def wire_body(json: BaseModel) -> dict[str, object]: + if isinstance(json, PartialBody): + return json.model_dump(by_alias=True, exclude_unset=True) + return json.model_dump(by_alias=True, exclude_none=True) + + def _headers(headers: BaseModel) -> dict[str, str]: dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} @@ -307,9 +319,26 @@ def request_with_retry[T: RetryableResponse]( return issue() -def _classify[R: BaseModel]( - resp: requests.Response, response_type: type[R] -) -> Result[R]: +class ClassifiableResponse(Protocol): + """What classifying an outcome reads off a response. requests.Response satisfies + it, and so does a fake, so the classification rules are testable on their own.""" + + @property + def status_code(self) -> int: ... + + @property + def ok(self) -> bool: ... + + @property + def text(self) -> str: ... + + @property + def content(self) -> bytes: ... + + def json(self) -> object: ... + + +def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]: if resp.status_code == 401: return UnauthorizedError(body=resp.text) if resp.status_code == 429: @@ -317,7 +346,8 @@ def _classify[R: BaseModel]( if not resp.ok: return UnknownApiError(status_code=resp.status_code, body=resp.text) try: - return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json())) + payload: Final[object] = resp.json() if resp.content else {} + return Success(status_code=resp.status_code, data=response_type.model_validate(payload)) except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value return ValidationError(message=str(exc)) @@ -335,13 +365,13 @@ def post[R: BaseModel]( lambda: requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get[R: BaseModel]( @@ -363,7 +393,7 @@ def get[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get_external[R: BaseModel]( @@ -383,7 +413,7 @@ def get_external[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def delete[R: BaseModel]( @@ -400,14 +430,14 @@ def delete[R: BaseModel]( lambda: requests.delete( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), params=_params(params), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def patch[R: BaseModel]( @@ -423,13 +453,13 @@ def patch[R: BaseModel]( lambda: requests.patch( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def put[R: BaseModel]( @@ -445,13 +475,13 @@ def put[R: BaseModel]( lambda: requests.put( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def probe( @@ -555,7 +585,7 @@ def send( str(url), headers=_headers(headers), params=_params(params), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=stream, timeout=timeout, ) @@ -605,7 +635,7 @@ def upload[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def stream_binary( @@ -623,7 +653,7 @@ def stream_binary( resp = requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=True, timeout=timeout, ) diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py index 9a45743a0cd..75817340876 100644 --- a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p from __future__ import annotations +from typing import Final, Literal + import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import StreamingResponse @@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel): class _BridgeChunk(BaseModel): id: str - choices: list[_BridgeChoice] = [] + object: Literal["chat.completion.chunk"] + choices: list[_BridgeChoice] = Field(default_factory=list) class _WeatherArgs(BaseModel): @@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming: resources.key(), ChatBody( model=bridged_model, - messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], max_tokens=64, stream=True, ), ) - chunks = _bridge_chunks(result) - ids = {chunk.id for chunk in chunks} + chunks: Final = _bridge_chunks(result) + assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk" + ids: Final = frozenset(chunk.id for chunk in chunks) assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" - assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id" @pytest.mark.covers( "llm.chat_completions.openai.basic.stream.bridge_streams_sse", @@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming: chunks = _bridge_chunks(result) content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" - assert any( - choice.finish_reason for chunk in chunks for choice in chunk.choices - ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), ( + f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + ) assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" @pytest.mark.covers( diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index d0f478185c2..368c20cb6aa 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations +import math +import random import time -from dataclasses import dataclass +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,17 +31,33 @@ from e2e_config import ( DD_SITE, POLL_TIMEOUT, ) -from e2e_http import URL, Headers, RateLimitedError, Success, post +from e2e_http import URL, Headers, StreamingResponse, send -#: How many rate-limited responses in a row one search tolerates before the -#: hard fail; each retry sleeps a full search interval, so this rides out a -#: burst from a concurrent consumer of the org-wide search budget. -_RATE_LIMIT_RETRIES = 5 +type SearchCall = Callable[[str, float], StreamingResponse] + + +def _seconds(value: str | None) -> float | None: + if value is None: + return None + try: + seconds: Final = float(value) + except ValueError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + +def _rate_limit_delay(headers: Mapping[str, str]) -> float: + delays: Final = tuple( + delay + for name in ("x-ratelimit-reset", "retry-after") + if (delay := _seconds(headers.get(name))) is not None + ) + return max(1.0, max(delays, default=DD_SEARCH_INTERVAL)) class _DdAuthHeaders(Headers): - api_key: str = Field(serialization_alias="DD-API-KEY") - app_key: str = Field(serialization_alias="DD-APPLICATION-KEY") + api_key: str = Field(serialization_alias="DD-API-KEY", repr=False) + app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False) class _SearchFilter(BaseModel): @@ -88,8 +108,12 @@ class _SearchResponse(BaseModel): @dataclass(frozen=True, slots=True) class DdLogsReader: site: str - api_key: str - app_key: str + api_key: str = field(repr=False) + app_key: str = field(repr=False) + search: SearchCall | None = field(default=None, repr=False) + now: Callable[[], float] = field(default=time.monotonic, repr=False) + sleep: Callable[[float], None] = field(default=time.sleep, repr=False) + jitter: Callable[[], float] = field(default=random.random, repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog @@ -108,25 +132,28 @@ class DdLogsReader: a single event. A 429 backs off and retries - the search budget is org-wide, so another consumer can empty it under us - while any other failure stays a hard fail.""" - for _ in range(_RATE_LIMIT_RETRIES): - result = post( - URL(f"https://api.{self.site}/api/v2/logs/events/search"), - headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=query)), - response_type=_SearchResponse, - timeout=30.0, - ) - match result: - case Success(data=page): - return [event.attributes for event in page.data] - case RateLimitedError(retry_after_seconds=retry_after): - time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL) - case failure: - pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}") + return self._events_for_query(query, self.now() + POLL_TIMEOUT) + + def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]: + search: Final = self.search or self._search_page + while (remaining := deadline - self.now()) > 0: + if (result := search(query, min(30.0, remaining))).ok: + return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data] + if result.status_code != 429: + pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}") + if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0: + self.sleep(delay) pytest.fail( - f"DataDog Logs Search API at api.{self.site} still rate-limited after " - f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide " - "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer" + f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; " + "the org-wide logs_public_search_api budget is exhausted" + ) + + def _search_page(self, query: str, timeout: float) -> StreamingResponse: + return send( + URL(f"https://api.{self.site}/api/v2/logs/events/search"), + headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), + json=_SearchRequest(filter=_SearchFilter(query=query)), + timeout=timeout, ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: @@ -140,33 +167,42 @@ class DdLogsReader: hide from the exactly-one assertion - real-DataDog jitter can surface one call's two events tens of seconds apart. Searches pace at DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's - request budget. At the deadline the last result is returned as-is.""" - deadline = time.monotonic() + POLL_TIMEOUT - while time.monotonic() < deadline: - events = self.events_for_query(query) + request budget. Discovery, quota retries, and duplicate detection share + one POLL_TIMEOUT deadline; an incomplete settle window fails closed.""" + deadline: Final = self.now() + POLL_TIMEOUT + while (remaining := deadline - self.now()) > 0: + events = self._events_for_query(query, deadline) if events: - return self._settled_events_for_query(query, events) - time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_query(query) + return self._settled_events_for_query(query, events, deadline) + if (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + return [] - def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. Keep the last non-empty result: a transient empty search (index lag) must not erase events already confirmed earlier in the settle window. + A successful final search must reach the full settle window before the + shared read-back deadline; otherwise duplicate detection is incomplete. """ - settle_deadline = time.monotonic() + DD_SETTLE_SECONDS + settle_deadline: Final = self.now() + DD_SETTLE_SECONDS last_nonempty = events - while time.monotonic() < settle_deadline: - time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_query(query) - if not latest: - continue + if len(events) > 1: + return events + while (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + if self.now() >= deadline: + break + latest = self._events_for_query(query, deadline) if len(latest) > 1: return latest - last_nonempty = latest - return last_nonempty + if latest: + last_nonempty = latest + if self.now() >= settle_deadline: + return last_nonempty + pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s") def build_dd_logs_reader() -> DdLogsReader: diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py new file mode 100644 index 00000000000..910a1cefd42 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,223 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Final + +import pytest + +from datadog_reader import DdLogsReader +from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization +from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT +from e2e_http import StreamingResponse + + +def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: + api_key: Final = "test-datadog-api-secret" + app_key: Final = "test-datadog-app-secret" + reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key) + headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key) + + for value in (reader, headers): + assert api_key not in repr(value) + assert app_key not in repr(value) + + assert headers.model_dump(by_alias=True) == { + "DD-API-KEY": api_key, + "DD-APPLICATION-KEY": app_key, + } + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +@dataclass +class Search: + responses: Iterator[StreamingResponse] + calls: tuple[tuple[str, float], ...] = () + + def __call__(self, query: str, timeout: float) -> StreamingResponse: + self.calls += ((query, timeout),) + return next(self.responses) + + +def _page(*event_ids: str) -> StreamingResponse: + return StreamingResponse( + status_code=200, + body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}), + ) + + +def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]: + search: Final = Search(iter(responses)) + return DdLogsReader( + site="us5.datadoghq.com", + api_key="test-api-secret", + app_key="test-app-secret", + search=search, + now=clock.now, + sleep=clock.sleep, + jitter=lambda: 0.25, + ), search + + +def test_429_honors_server_reset_and_preserves_duplicate_events() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")), + clock, + ) + + events: Final = reader.events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 6.25 + assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0)) + + +@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1")) +def test_invalid_reset_uses_search_interval(reset: str) -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25 + + +def test_zero_reset_cannot_create_a_busy_retry_loop() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 1.25 + + +def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 8.25 + + +def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls == (("test-marker", 30.0),) + + +def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75)) + + +@pytest.mark.parametrize("status", (-1, 401, 403, 500)) +def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None: + clock: Final = Clock() + reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock) + + with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"): + reader.events_for_query("test-marker") + + assert search.calls == (("test-marker", 30.0),) + assert clock.elapsed == 0 + + +def test_polling_quota_retries_share_the_original_deadline() -> None: + clock: Final = Clock() + reader, search = _reader( + (_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == 2 + + +def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None: + clock: Final = Clock() + attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) + reader, search = _reader((_page(),) * attempts, clock) + + assert reader.poll_events_for_query("test-marker") == [] + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == attempts + + +def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader( + (_page(),) * empty_reads + + (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL) + assert len(search.calls) == empty_reads + 2 + + +def test_settlement_detects_a_duplicate_on_the_final_search() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 30 + + +def test_settlement_keeps_confirmed_events_through_empty_searches() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first",) + assert clock.elapsed == 30 + + +def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock) + + with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == empty_reads + 2 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 2b897f5f07f..1ef0d89a8f9 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -43,6 +43,9 @@ from models import ( KeyResetSpendBody, KeyResetSpendResponse, KeyUpdateBody, + McpServerCreateBody, + McpServerRow, + McpServerUpdateBody, ModelDeleteBody, OrgDeleteBody, OrgInfoParams, @@ -537,6 +540,38 @@ class ManagementClient: ).root ) + def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow: + return unwrap( + self.proxy.transport.post( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow: + """PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial + update where a field left unset keeps its stored value and None clears it.""" + return unwrap( + self.proxy.transport.put( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def delete_mcp_server(self, server_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can + unwrap it while a deferred teardown can ignore an already-deleted server.""" + return self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", diff --git a/tests/e2e/management/test_mcp_lifecycle_e2e.py b/tests/e2e/management/test_mcp_lifecycle_e2e.py new file mode 100644 index 00000000000..9257d697647 --- /dev/null +++ b/tests/e2e/management/test_mcp_lifecycle_e2e.py @@ -0,0 +1,294 @@ +"""Live e2e: the MCP server and toolset management routes' lifecycle contract. + +Two customer defects sit on these routes, and each step here is the read-back that +would have caught one of them: a dashboard edit that took several saves to stick +because the read landed on a replica the write had not reached, and a toolset whose +tools were stored under one name and read back under another, so it granted +nothing. Every read-back therefore polls every replica that serves the route +(ProxyClient.read_back_everywhere) and asserts the exact values written, and both +update routes are held to the same partial-update contract: a field left out of the +payload keeps its stored value, a field sent as null is cleared. The server URL is +unreachable on purpose; only persistence is under test, never a tool call. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + McpInfo, + McpServerCreateBody, + McpServerListResponse, + McpServerRow, + McpServerUpdateBody, + ToolsetCreateBody, + ToolsetListResponse, + ToolsetRow, + ToolsetTool, + ToolsetUpdateBody, +) + +pytestmark = pytest.mark.e2e + +UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp" + + +def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]: + name: Final = f"e2e_mcp_lifecycle_{unique_marker()}" + body: Final = McpServerCreateBody( + server_name=name, + alias=name, + url=UNREACHABLE_URL, + transport="http", + description="e2e lifecycle server", + mcp_info=McpInfo( + server_name=f"{name} (display)", + description="shown on the MCP page", + logo_url="https://e2e.test.local/logo.png", + ), + ) + server_id: Final = client.create_mcp_server(body).server_id + resources.defer(lambda: client.delete_mcp_server(server_id)) + return body, server_id + + +def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None: + stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info) + expected: Final = ( + written.server_name, + written.alias, + written.url, + written.transport, + written.description, + written.mcp_info, + ) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _server_everywhere( + client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool] +) -> Mapping[str, McpServerRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled) + + +def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]: + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: any(row.server_id == server_id for row in rows.root), + ) + return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()} + + +class TestMcpServerLifecycle: + @pytest.mark.covers("mgmt.mcp_server.new.persists") + def test_create_persists_every_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id) + for replica, row in by_id.items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}") + + @pytest.mark.skip( + reason=( + "product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose " + "_build_mcp_server_table sets description from mcp_info['description'], so the list " + "reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the " + "stored description column. A server created with both set to different text reads " + "back with two different descriptions depending on the route" + ) + ) + @pytest.mark.covers("mgmt.mcp_server.list.persists") + def test_created_server_is_listed_with_every_field( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + for replica, row in _listed_server_everywhere(client, server_id).items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}") + + @pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields") + def test_updating_only_the_alias_keeps_every_other_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + renamed: Final = f"{body.alias}_renamed" + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed)) + + after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed) + for replica, row in after_one_put.items(): + _assert_server_matches( + row, + body.model_copy(update={"alias": renamed}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias", + ) + + @pytest.mark.covers("mgmt.mcp_server.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None)) + + cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_server_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_server.delete.persists") + def test_delete_removes_the_server_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + + _ = unwrap(client.delete_mcp_server(server_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}") + assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: all(row.server_id != server_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.server_id != server_id for row in rows.root), ( + f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}" + ) + + +def _create_toolset( + client: ManagementClient, resources: ResourceManager, server_id: str +) -> tuple[ToolsetCreateBody, str]: + body: Final = ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="e2e lifecycle toolset", + tools=[ + ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"), + ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"), + ], + ) + toolset_id: Final = client.proxy.create_toolset(body).toolset_id + resources.defer(lambda: client.proxy.delete_toolset(toolset_id)) + return body, toolset_id + + +def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None: + stored: Final = (row.toolset_name, row.description, row.tools) + expected: Final = (written.toolset_name, written.description, written.tools) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _toolset_everywhere( + client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool] +) -> Mapping[str, ToolsetRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled) + + +class TestMcpToolsetLifecycle: + @pytest.mark.covers("mgmt.mcp_toolset.new.persists") + def test_create_persists_both_tools_under_the_exact_names_written( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id) + for replica, row in by_id.items(): + _assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}") + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + _assert_toolset_matches( + next(row for row in rows.root if row.toolset_id == toolset_id), + body, + where=f"GET /v1/mcp/toolset on {replica}", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields") + def test_updating_only_the_description_keeps_the_tools_and_name( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited")) + + edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited") + for replica, row in edited.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": "edited"}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.persists") + def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + kept: Final = body.tools[:1] + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept)) + + narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept) + for replica, row in narrowed.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"tools": kept}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None)) + + cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.delete.persists") + def test_delete_removes_the_toolset_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + _, toolset_id = _create_toolset(client, resources, server_id) + + _ = unwrap(client.proxy.delete_toolset(toolset_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}") + assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.toolset_id != toolset_id for row in rows.root), ( + f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}" + ) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..352b4446cfd 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the core Datadog toolset with its credentials from the env. By default + the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose + every tool the core toolset serves.""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=None if allowed_tools is None else list(allowed_tools), mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..210fc7a1e98 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -16,11 +16,11 @@ import time from collections.abc import Mapping from dataclasses import dataclass -from pydantic import BaseModel, ConfigDict, Field, RootModel +from pydantic import BaseModel, ConfigDict, Field from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel): server_id: str -class McpServerRow(BaseModel): - server_id: str - alias: str | None = None - url: str | None = None - - -class McpServersListResponse(RootModel[list[McpServerRow]]): - pass - - class McpToolMcpInfo(BaseModel): server_id: str | None = None alias: str | None = None @@ -193,7 +183,7 @@ class McpClient: "/v1/mcp/server", headers=self.proxy.transport.master, params=NoBody(), - response_type=McpServersListResponse, + response_type=McpServerListResponse, ) ).root @@ -224,11 +214,16 @@ class McpClient: user_id: str, mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, + mcp_toolsets: list[str] | None = None, models: list[str] | None = None, ) -> str: object_permission = ( - ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) - if mcp_servers is not None or mcp_access_groups is not None + ObjectPermission( + mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, + mcp_toolsets=mcp_toolsets, + ) + if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None else None ) return self.proxy.generate_key( @@ -272,6 +267,20 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]: + """Poll tools/list until `server_id`'s tools as `key` sees them are exactly + `expected`, and return the last listing either way, so the caller's equality + assertion names the difference. Fails at poll_timeout only when the read + itself never succeeded.""" + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected: + return expected + if time.monotonic() >= deadline: + return unwrap(result).tool_names_for_server(server_id) + time.sleep(self.proxy.poll_interval) + def await_call_tool( self, key: str, diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..6b901145eb1 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,95 @@ +"""Live e2e: a key granted a toolset lists exactly the toolset's tools. + +An admin registers the real Datadog remote MCP server with its whole core toolset +exposed, discovers two of its tool names through a key granted the server outright, +and curates a toolset naming exactly those two. A second key is granted the server +plus that toolset, and its tools/list must come back as exactly those two names: no +more, so the rest of the server's catalog stays hidden behind the toolset, and no +fewer, so a tool stored under one name and read under another (which granted +nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP +upstream). +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ToolsetCreateBody, ToolsetTool + +pytestmark = pytest.mark.e2e + + +def _key( + client: McpClient, + resources: ResourceManager, + label: str, + *, + server_id: str, + toolset_id: str | None = None, +) -> str: + key: Final = client.generate_key( + user_id=f"e2e-mcp-{label}-{unique_marker()}", + mcp_servers=[server_id], + mcp_toolsets=None if toolset_id is None else [toolset_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _wire_prefix(wire_name: str, tool_name: str, catalog: frozenset[str]) -> str: + """The prefix tools/list puts in front of one server's tool names, measured off a + tool whose own name is known rather than guessed from the alias. A toolset grants + by the tool's own name, never the wire name, and the prefix is whatever the proxy + is configured to build (the alias, or a short server id), so measuring it is the + only way to cross between the two.""" + assert wire_name.endswith(tool_name), f"tools/list served {wire_name!r}, expected it to end with {tool_name!r}" + prefix: Final = wire_name[: len(wire_name) - len(tool_name)] + unprefixed: Final = frozenset(name for name in catalog if not name.startswith(prefix)) + assert not unprefixed, ( + f"every tool of one server shares the wire prefix {prefix!r}, so {sorted(unprefixed)} " + f"cannot be reduced to the names a toolset grants by" + ) + return prefix + + +class TestMcpToolsetEnforcement: + @pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped") + def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None: + server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + + catalog_key: Final = _key(client, resources, "catalog", server_id=server_id) + known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL) + catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id) + assert len(catalog) > 2, ( + f"the Datadog core toolset must serve more tools than the toolset names, or the " + f"restriction has nothing to hide; got {sorted(catalog)}" + ) + prefix: Final = _wire_prefix(known_wire, SEARCH_LOGS_TOOL, catalog) + chosen_wire: Final = frozenset(sorted(catalog)[:2]) + chosen: Final = frozenset(name.removeprefix(prefix) for name in chosen_wire) + + toolset: Final = client.proxy.create_toolset( + ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="two Datadog tools", + tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)], + ) + ) + resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id)) + assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, ( + f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim" + ) + + scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id) + listed: Final = client.await_tools(scoped_key, server_id, expected=chosen_wire) + assert listed == chosen_wire, ( + f"a key granted the toolset must list exactly its two tools; " + f"got {sorted(listed)}, expected {sorted(chosen_wire)}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f2654e0eec..62810e6cfd9 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -10,6 +10,7 @@ from collections.abc import Sequence from datetime import datetime from typing import Final, Literal +from e2e_http import PartialBody from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -55,6 +56,7 @@ class KeyMetadata(BaseModel): class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -77,7 +79,7 @@ class KeyGenerateBody(BaseModel): allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None - router_settings: "RouterSettingsOverride | None" = None + router_settings: RouterSettingsOverride | None = None class KeyGenerateResponse(BaseModel): @@ -516,6 +518,15 @@ class CountTokensResponse(BaseModel): # ---------- mcp servers ---------- +class McpInfo(BaseModel): + """The `mcp_info` display block stored on an MCP server; only the fields the + lifecycle test writes and reads back.""" + + server_name: str | None = None + description: str | None = None + logo_url: str | None = None + + class McpServerCreateBody(BaseModel): """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints @@ -530,6 +541,18 @@ class McpServerCreateBody(BaseModel): oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None authorization_url: str | None = None token_url: str | None = None + server_name: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerUpdateBody(PartialBody): + """PUT /v1/mcp/server: a field left unset keeps its stored value, a field set + to None is cleared.""" + + server_id: str + alias: str | None = None + description: str | None = None class McpServerInfo(BaseModel): @@ -543,6 +566,54 @@ class McpServerInfo(BaseModel): allow_all_keys: bool | None = None +class McpServerRow(McpServerInfo): + """A stored MCP server as the create, get, and list routes return it: the + fields the lifecycle test asserts survive the round trip.""" + + server_name: str | None = None + transport: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerListResponse(RootModel[list[McpServerRow]]): + """GET /v1/mcp/server answers with a bare array of servers.""" + + +class ToolsetTool(BaseModel): + server_id: str + tool_name: str + + +class ToolsetCreateBody(BaseModel): + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] + + +class ToolsetUpdateBody(PartialBody): + """PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set + to None is cleared.""" + + toolset_id: str + description: str | None = None + tools: list[ToolsetTool] | None = None + + +class ToolsetRow(BaseModel): + """A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id}, + and each row of GET /v1/mcp/toolset return it.""" + + toolset_id: str + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] = Field(default_factory=list) + + +class ToolsetListResponse(RootModel[list[ToolsetRow]]): + """GET /v1/mcp/toolset answers with a bare array of toolsets.""" + + class EmbedBody(BaseModel): model: str input: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 520cbfde5a9..1bac5116a9d 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable, Mapping from dataclasses import dataclass +from functools import reduce from datetime import datetime from types import MappingProxyType from typing import Final @@ -26,6 +27,7 @@ from e2e_http import ( Result, StreamingResponse, Success, + UnknownApiError, is_ok, unwrap, ) @@ -70,6 +72,9 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + ToolsetCreateBody, + ToolsetRow, + ToolsetUpdateBody, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -82,7 +87,7 @@ from e2e_config import ( SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, ) -from transport import HttpTransport, SplitTransport, Transport +from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -235,6 +240,99 @@ def servable_timeout_message( ) +type ReplicaRead[T] = Callable[[float], T] + + +@dataclass(frozen=True, slots=True) +class EverywhereConverged[T]: + """Every replica answered with something `settled` accepts, keyed by replica.""" + + answers: Mapping[str, T] + + +@dataclass(frozen=True, slots=True) +class NeverConvergedOn[T]: + """`replica` ran out its budget without an answer `settled` accepts; `last` is + its final answer, so the failure can say what that replica still serves.""" + + replica: str + last: T + + +def _last_answer[T]( + read: ReplicaRead[T], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> T: + """Poll `read` until `settled` accepts its answer or `timeout` runs out, and + return the last answer either way. Each read's request timeout is clamped to + the budget left, and the final poll runs even when less than an interval + remains, so a deadline never skips the read that would have settled.""" + deadline: Final = now() + timeout + answer = read(min(request_timeout, timeout)) + while not settled(answer): + remaining = deadline - now() + if remaining <= 0: + return answer + sleep(min(interval, remaining)) + answer = read(min(request_timeout, remaining)) + return answer + + +def await_everywhere[T]( + reads: Mapping[str, ReplicaRead[T]], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> EverywhereConverged[T] | NeverConvergedOn[T]: + """`_last_answer` against every replica in turn, each with the full budget, so a + write counts as visible only once the last replica reflects it, and stop at the + first replica that never converges. Clock and sleep are injected.""" + def read_replica( + outcome: EverywhereConverged[T] | NeverConvergedOn[T], + item: tuple[str, ReplicaRead[T]], + ) -> EverywhereConverged[T] | NeverConvergedOn[T]: + if isinstance(outcome, NeverConvergedOn): + return outcome + replica, read = item + answer: Final = _last_answer( + read, + settled=settled, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ) + if not settled(answer): + return NeverConvergedOn(replica=replica, last=answer) + return EverywhereConverged(answers=MappingProxyType({**outcome.answers, replica: answer})) + + initial: Final[EverywhereConverged[T] | NeverConvergedOn[T]] = EverywhereConverged(answers=MappingProxyType({})) + return reduce(read_replica, reads.items(), initial) + + +def _is_not_found[R: BaseModel](result: Result[R]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _status_of[R: BaseModel](result: Result[R]) -> int: + match result: + case Success(status_code=status_code) | UnknownApiError(status_code=status_code): + return status_code + case _: + return -1 + + type Poller[T] = Callable[[], T] @@ -321,6 +419,7 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re class ProxyClient: transport: Transport replicas: Mapping[str, Transport] + control_replicas: Mapping[str, Transport] poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -569,6 +668,112 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2) + # ---- replica read-back ---------------------------------------------- + + def replicas_for(self, path: str) -> Mapping[str, Transport]: + """The replicas that serve `path`: every data-plane replica for an LLM route, + and for a management route the control-plane replicas, since the data-plane + replicas trim management routes and answer them 404. A monolith serves both + from every replica, so a management read-back polls all of them; a split + deployment exposes one control-plane address (there is one backend process + behind it on the stack these suites run against), so it polls that. A + control plane fronting several backends would need its own replica list to + prove each one converged, the way PROXY_REPLICA_URLS does for the gateways. + Never empty: a read-back against no replica would assert nothing and pass.""" + replicas: Final = self.control_replicas if is_control_plane_path(path) else self.replicas + assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing" + return replicas + + def read_body_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, settled: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica that serves it, polling each to poll_timeout + until `settled` accepts its body, and fail naming the first replica that + never converged. Returns each replica's settled body, keyed by replica, so + the caller can assert the rest of it.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()}, + settled=lambda result: isinstance(result, Success) and settled(result.data), + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: unwrap(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; " + f"last read: {last}" + ) + + def gone_everywhere(self, path: str) -> Mapping[str, int]: + """Poll GET `path` on every replica that serves it until each stops serving + it, and fail naming the first replica that still does at poll_timeout. + Returns each replica's final status, so the caller asserts the 404 itself.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()}, + settled=_is_not_found, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: _status_of(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" + ) + + @staticmethod + def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + return lambda request_timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=request_timeout, + ) + + # ---- mcp toolsets --------------------------------------------------- + + def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow: + return unwrap( + self.transport.post( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow: + """PUT /v1/mcp/toolset: a partial update where a field left unset keeps its + stored value and None clears it.""" + return unwrap( + self.transport.put( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def delete_toolset(self, toolset_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase + can unwrap it while a deferred teardown can ignore an already-deleted row.""" + return self.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def create_credential(self, body: CredentialCreateBody) -> None: unwrap( self.transport.post( @@ -736,7 +941,10 @@ def build_proxy_client( base URLs are the same for a monolithic proxy, so routing is then a no-op. ``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model barrier polls directly; it is the data-plane URL itself unless the stack - exports each gateway's own address. + exports each gateway's own address. Management read-backs poll those same + replicas when the two planes share a base URL (a monolith, where every replica + serves every route) and the control plane alone when they differ (a split + deployment, where the data-plane replicas do not serve management routes). The endpoints are injectable for callers that resolve the proxy some other way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must @@ -764,9 +972,13 @@ def build_proxy_client( for url in replica_urls } ) + control_replicas: Final = ( + replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control}) + ) return ProxyClient( transport=split, replicas=replicas, + control_replicas=control_replicas, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 188db2a8eb5..374badcf5fc 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form. import json import os from collections.abc import Iterator +from contextlib import ExitStack from dataclasses import dataclass from typing import Final @@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel): @dataclass(frozen=True, slots=True) -class TagSplitDeployments: - """Scenario A mirrors the customer-shaped config from GitHub issue #36619: - plain deployment registered first, tier deployment and marker both tagged. - Scenario B flips both axes for GitHub issue #36621: marker registered first - and its tier deployment left untagged, so routing depends neither on - registration order nor on tier deployments carrying tags.""" - - tag_a: str - shared_a: str - tier_a: str - tag_b: str - shared_b: str - tier_b: str +class TagSplitDeployment: + tag: str + shared: str + tier: str @dataclass(frozen=True, slots=True) @@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]: } -def _key_for( - proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False -) -> str: +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str: key: Final = proxy.generate_key( KeyGenerateBody( models=models, @@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con ) -@pytest.fixture(scope="module") -def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _register_models( + proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...] +) -> None: + for name, params in registrations: + stack.callback(proxy.delete_model, proxy.create_model(name, params)) + + +def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment: marker: Final = unique_marker() - deployments: Final = TagSplitDeployments( - tag_a=f"e2e-split-a-{marker}", - shared_a=f"e2e-autoroute-a-{marker}", - tier_a=f"e2e-tier-a-{marker}", - tag_b=f"e2e-split-b-{marker}", - shared_b=f"e2e-autoroute-b-{marker}", - tier_b=f"e2e-tier-b-{marker}", + named: Final = TagSplitDeployment( + tag=f"e2e-split-{marker}", + shared=f"e2e-autoroute-{marker}", + tier=f"e2e-tier-{marker}", ) anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") - marker_params_a: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_a), - tags=[deployments.tag_a], + marker_registration: Final = ( + named.shared, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + tags=[named.tag], + ), ) - marker_params_b: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_b), - tags=[deployments.tag_b], + tier_registration: Final = ( + named.tier, + LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]), ) - registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( - (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), - (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), - (deployments.shared_a, marker_params_a), - (deployments.shared_b, marker_params_b), - (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), - (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)) + registrations: Final = ( + (marker_registration, tier_registration, plain_registration) + if marker_first + else (plain_registration, tier_registration, marker_registration) ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield deployments - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, stack, registrations) + return named -@pytest.fixture(scope="module") -def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: +@pytest.fixture(scope="class") +def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=False) + + +@pytest.fixture(scope="class") +def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=True) + + +@pytest.fixture(scope="class") +def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias: marker: Final = unique_marker() named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: +@pytest.fixture(scope="class") +def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit: marker: Final = unique_marker() named: Final = HeuristicSplit( alias=f"e2e-heuristic-router-{marker}", @@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: +@pytest.fixture(scope="class") +def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter: marker: Final = unique_marker() named: Final = SemanticAutoRouter( marker=f"e2e-semantic-router-{marker}", @@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.marker, marker_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: +@pytest.fixture(scope="class") +def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias: marker: Final = unique_marker() named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named class TestTagSplitRouting: @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") def test_body_tagged_chat_routes_through_the_marker_to_its_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36619: with tag filtering on, a chat request whose body metadata tags match the tagged marker under a shared model name is answered by the marker's tier deployment, not by the plain deployment that was registered under the name first.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag]))) assert chat.choices, "tagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_chat_is_always_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620: untagged chat requests to the shared name succeed on every call and are all served by the plain deployment; the tagged marker never captures them, so no intermittent auto-router errors and no tier hijacking.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) for _ in range(5): - chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared))) assert chat.choices, "untagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=5) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + _assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_messages_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620 on the /v1/messages surface: an untagged Anthropic-native request to the shared name is served by the plain deployment, not captured by the tagged marker.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared))) assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name" + ) class TestUntaggedTierDeployments: @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36621: a /v1/messages request tagged only via the x-litellm-tags header selects the tagged marker, and the rewrite still lands on the tier deployment even though that deployment carries no tags, because the marker consumed the routing tags.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag) answer: Final = unwrap( proxy.transport.post( "/v1/messages", headers=headers, - json=_hello_messages_body(split.shared_b), + json=_hello_messages_body(marker_first_split.shared), response_type=AnthropicMessagesResponse, ) ) assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + _assert_served_only_by( + rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name" + ) @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins the tag-consumption half of GitHub issue #36621: after the tagged marker rewrites the request to its tier model, the consumed routing tags no longer constrain deployment selection, so the untagged tier deployment serves the request instead of a strict-tag denial.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + chat: Final = unwrap( + proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag])) + ) assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + _assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier") @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """The tag-consumption fix must not loosen strict tag semantics: a tagged request aimed directly at an untagged deployment (no marker involved) is still rejected with the 401 tags-configuration error.""" - key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) - result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag])) assert isinstance(result, UnauthorizedError), ( f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" ) @@ -451,37 +450,39 @@ class TestUntaggedTierDeployments: class TestResponsesApiTagRouting: @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_header_tagged_responses_with_string_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the /v1/responses surface of the tag split (GitHub issues #36620/#36621): a /v1/responses request with string input, tagged via the x-litellm-tags header, succeeds and routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) ) assert answer.id, "header-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + _assert_served_only_by( + rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input" + ) @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_body_tagged_responses_with_list_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the body-tag and list-input combination of the same split: /v1/responses with litellm_metadata.tags and structured input items routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, + model=plain_first_split.shared, input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], max_output_tokens=64, - litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]), ) answer: Final = unwrap( proxy.transport.post( @@ -493,18 +494,18 @@ class TestResponsesApiTagRouting: ) assert answer.id, "body-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_responses_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the untagged half of the /v1/responses tag split: an untagged request to the shared name is served by the plain deployment, matching the chat and messages surfaces.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post( @@ -516,7 +517,9 @@ class TestResponsesApiTagRouting: ) assert answer.id, "untagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name" + ) class TestStrategyAliasPricing: @@ -551,9 +554,7 @@ class TestComplexityHeuristicScope: while the accompanying ~2KB agent system prompt is packed with enough reasoning and complexity keywords that scoring the combined text lands in REASONING; only ask-only scoring keeps this on the cheap tier.""" - key: Final = _key_for( - proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] - ) + key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]) body: Final = ChatBody( model=heuristic_split.alias, messages=[ diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 66841725d1d..81cd6c8d3d1 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -13,13 +13,24 @@ monkeypatches anything. from __future__ import annotations from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from types import MappingProxyType from typing import Final import pytest - -from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome +from e2e_http import ( + RETRY_ATTEMPTS, + TRANSIENT_STATUSES, + NoBody, + PartialBody, + Success, + ValidationError, + classify, + request_with_retry, + streaming_outcome, + wire_body, +) +from pydantic import BaseModel, TypeAdapter @dataclass @@ -33,10 +44,10 @@ class FakeResponse: @dataclass class SleepRecorder: - delays: list[float] = field(default_factory=list) + delays: tuple[float, ...] = () def __call__(self, seconds: float) -> None: - self.delays.append(seconds) + self.delays += (seconds,) def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]: @@ -55,7 +66,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_429_is_never_retried(self) -> None: @@ -63,7 +74,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_overloaded_529_retries_with_backoff_then_returns_the_success(self) -> None: @@ -71,7 +82,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[1] - assert sleep.delays == [0.5] + assert sleep.delays == (0.5,) assert responses[0].close_calls == 1 assert responses[1].close_calls == 0 @@ -80,7 +91,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[RETRY_ATTEMPTS - 1] - assert sleep.delays == [0.5, 1.0] + assert sleep.delays == (0.5, 1.0) assert [r.close_calls for r in responses] == [1, 1, 0, 0] @@ -134,3 +145,65 @@ class TestStreamEventArrivals: assert result.stream_events == [] assert result.stream_event_arrivals == [] assert result.body == "bad request" + + +class _ServerUpdate(PartialBody): + server_id: str + alias: str | None = None + description: str | None = None + + +class _ServerCreate(BaseModel): + alias: str + description: str | None = None + + +class TestWireBody: + """A partial-update body must put exactly the caller's choice on the wire: an + omitted field stays off it so the route keeps the stored value, and an explicit + None goes out as JSON null so the route clears it. Plain bodies keep dropping + None, which is what every create route expects.""" + + def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None: + assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None} + assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"} + + def test_plain_body_drops_none_fields(self) -> None: + assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"} + + +_JSON: Final[TypeAdapter[object]] = TypeAdapter(object) + + +@dataclass +class FakeJsonResponse: + """The `classify` view of a response: a status, the raw body bytes, and the + parse that would raise on an empty one.""" + + status_code: int + content: bytes + + @property + def ok(self) -> bool: + return self.status_code < 400 + + @property + def text(self) -> str: + return self.content.decode() + + def json(self) -> object: + return _JSON.validate_json(self.content) + + +class TestClassifyEmptyBody: + """A delete that answers 202 with no body is a success, not a parse failure: + the MCP server and toolset delete routes both answer that way, and reading it + as a failure would hide a delete that did not happen behind one that did.""" + + def test_empty_2xx_body_is_a_success(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody) + assert isinstance(result, Success) and result.status_code == 202 + + def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=200, content=b""), NoBody) + assert isinstance(result, ValidationError) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 2caac58333f..3b84a47e3cc 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -15,28 +15,35 @@ from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat from types import MappingProxyType -from typing import Final +from typing import Final, cast import pytest - from e2e_config import parse_replica_urls from e2e_http import Result, Success from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( - Poller, ConvergeOutcome, Converged, + EverywhereConverged, ModelsPoller, + NeverConvergedOn, NotConverged, NotServableOn, + Poller, + ProxyClient, + ReplicaRead, Servable, await_converged_everywhere, + await_everywhere, await_servable_everywhere, - first_lagging_replica, + build_proxy_client, converge_timeout_message, + first_lagging_replica, ) +from transport import Transport MODEL: Final = "gpt-under-test" +_NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 RPM_BEFORE_UPDATE: Final = 100 @@ -187,3 +194,83 @@ class TestParseReplicaUrls: def test_falls_back_to_the_data_plane_address_when_unset(self) -> None: assert parse_replica_urls("", "http://lb") == ("http://lb",) + + +def _answers(answers: Iterable[str]) -> ReplicaRead[str]: + it: Final = iter(answers) + return lambda _timeout: next(it) + + +def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> EverywhereConverged[str] | NeverConvergedOn[str]: + clock: Final = FakeClock() + return await_everywhere( + reads, + settled=lambda answer: answer == "renamed", + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))), + } + outcome: Final = _await_everywhere(reads) + assert isinstance(outcome, EverywhereConverged) + assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"} + + def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(repeat("stale")), + } + assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale") + + def test_polls_until_the_deadline_before_giving_up(self) -> None: + lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed")) + outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)}) + assert isinstance(outcome, EverywhereConverged), outcome + + +class TestReplicasFor: + def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} + + def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://lb", + replica_urls=("http://pod-1", "http://pod-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://pod-1", "http://pod-2"} + + def test_mcp_admin_routes_read_back_from_every_data_plane_replica(self) -> None: + """/v1/mcp/* is a lazily mounted feature, so a data-plane replica serves it + too and answers from its own in-memory registry. Routing it to the control + plane would leave every replica but that one unproven, and would move the + tools/list barrier in mcp_client off the plane that serves tools/list.""" + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://gateway-1", "http://gateway-2"} + assert set(client.replicas_for("/v1/mcp/toolset/abc")) == {"http://gateway-1", "http://gateway-2"} + + def test_a_route_no_replica_serves_is_refused_rather_than_read_back_vacuously(self) -> None: + """A read-back over zero replicas would satisfy every predicate and assert + nothing, so asking for one fails instead of passing silently.""" + client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) + with pytest.raises(AssertionError, match="no replica is configured"): + _ = client.replicas_for("/v1/models") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index c0d055edb7f..669e094fee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -1,10 +1,11 @@ """ -Tests for partial-update semantics of PUT /v1/mcp/server. +Tests for partial-update semantics of PUT /v1/mcp/server and PUT /v1/mcp/toolset. A partial update must only write the fields the caller explicitly provided. Omitting a field must NOT reset it to its Pydantic schema default (e.g. ``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which -would silently overwrite the existing DB row. +would silently overwrite the existing DB row, and a field the caller sent as null +must be cleared rather than left at its stored value. """ import json @@ -850,3 +851,69 @@ async def test_cf_pair_switch_does_not_clear_dcr_bridge(): data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") assert "dcr_bridge" not in data_dict + + +def _mock_toolset_prisma(): + """A prisma double whose update answers with a row the reader can expand, so the + call under test returns instead of failing inside the row mapper.""" + updated_row = MagicMock() + updated_row.model_dump.return_value = { + "toolset_id": "ts-1", + "toolset_name": "ops", + "description": None, + "tools": "[]", + } + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcptoolsettable = AsyncMock() + mock_prisma.db.litellm_mcptoolsettable.update = AsyncMock(return_value=updated_row) + return mock_prisma + + +async def _run_toolset_update(payload: dict) -> dict: + """The columns PUT /v1/mcp/toolset writes for this payload, minus the audit stamp + every write carries. The prisma double is injected, so nothing is patched.""" + from litellm.proxy._experimental.mcp_server.toolset_db import update_mcp_toolset + from litellm.types.mcp_server.mcp_toolset import UpdateMCPToolsetRequest + + mock_prisma = _mock_toolset_prisma() + await update_mcp_toolset(mock_prisma, UpdateMCPToolsetRequest.model_validate(payload), "test-user") + written = dict(mock_prisma.db.litellm_mcptoolsettable.update.call_args[1]["data"]) + assert written["updated_by"] == "test-user" + return {name: value for name, value in written.items() if name != "updated_by"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_clears_description_on_explicit_null(): + """The dump used to drop None, so a null description could never clear the stored + one: the toolset kept a description its owner had deleted.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "description": None}) == {"description": None} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_omits_the_fields_the_caller_left_out(): + tools = [{"server_id": "s1", "tool_name": "alpha"}] + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": tools}) == {"tools": json.dumps(tools)} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_null_tools_rather_than_revoking_them(): + """A client that sends tools=null means "leave the selection alone", so the grants + survive. Clearing them is an explicit [], which cannot be confused with an omitted + field; treating null as a clear would silently revoke every tool the toolset grants.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": None, "description": "kept"}) == { + "description": "kept" + } + + +@pytest.mark.asyncio +async def test_toolset_partial_update_empties_the_selection_on_an_explicit_empty_list(): + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": []}) == {"tools": "[]"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_a_null_name(): + """A toolset always has a name, so a null toolset_name is a no-op, not a clear + that would write a NOT NULL column to null.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "toolset_name": None, "description": "kept"}) == { + "description": "kept" + } From e8e3172d7d70558929f32f057ebe4c7471c8c352 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 23:10:18 -0700 Subject: [PATCH 087/157] fix(model-management): honor an explicit null as a clear on model update (#40047) * fix(model-management): honor an explicit null as a clear on model update PATCH /model/{model_id}/update merged the patch with exclude_none and then popped explicit nulls only for the mirrored pricing fields, so a null sent for max_input_tokens, mode, supports_vision or any other key was dropped and a value pinned by an earlier save could never be removed. The route now follows JSON Merge Patch over both blobs: a key absent from the body is unchanged, a key sent as null is removed from the stored row, and a key sent with a value is set. Ownership and identity keys keep ignoring a null, as do the fields the stored models require, since clearing one writes a row no reload can rebuild. Mirrored pricing keys still clear from both blobs. Clearing a price also needed the router to stop merging a deployment's cost-map entry onto its previous registration, which left the old rate in place and kept billing at a price the deployment no longer carried. Adds a create, read, partial-update, clear, enforce, delete lifecycle e2e that reads back on every replica, and a harness helper for that read-back. * fix(router): keep a deployment id that names a real model from evicting its catalog entry Deployments are keyed into litellm.model_cost alongside the built-in catalog, so evicting a deployment's stale entry by id could take a real model's entry with it: registering a deployment whose model_info.id is "gpt-4o" stripped that model's pricing, context window and capability flags process-wide, for every other deployment of it, until the next price-map reload. Only evict an entry this registration owns. A colliding id keeps the previous merge, which pollutes the catalog entry rather than emptying it. Also pins the Admin UI round trip: the model edit form echoes the whole /model/info row back on save, and that read reports every key the deployment never stored as an explicit null, so the clear path has to leave those keys alone. * fix(router): decide cost-map eviction by what this registrar created The previous guard read a catalog entry off `litellm_provider`, so a deployment that declares its own provider in model_info was treated as one and kept billing at a price it no longer carried. It also only held for a single registration: a second one under a colliding id saw the id the first merge left behind and evicted the catalog entry anyway. Track the cost-map keys this registrar creates instead. A key it created is evicted before re-registration; one it did not is left to merge, which is what a deployment id colliding with a catalog model name needs. Also folds the required-fields comment into the docstring that already gives the reason. * fix(router): release a deployment's cost-map key when it is deleted The ownership ledger only grew. A deleted deployment kept its claim, so if a later catalog refresh started publishing a model under that same name, the next registration would treat the catalog entry as the deployment's own and evict it. Deleting a deployment now gives the key back, which also stops the ledger growing for the life of the process. * fix(router): hold a cost-map key while another live router still serves it The claim is process-wide but the release was per-deletion, so with two routers serving one deployment id, the first deletion put the survivor back on merging and the price it had just cleared would keep billing. Release the key only once no live router still serves that id. * fix(router): register a router in the live set when it gains a deployment _live_routers was only joined when a router was constructed with a model_list, but a router built empty is populated through add_deployment, and the empty branch exists for exactly that. Such a router was invisible to the live-router scan, so deleting the deployment from another router released the shared cost-map key while it was still serving that id. Joining the set where a deployment enters the list covers every path, and it also lets a price reload rebuild what a dynamically built router serves. * fix(e2e): read the stored model row from the control plane, not each gateway The lifecycle suite polled /model/info on every URL in PROXY_REPLICA_URLS. Those URLs are the stack's gateways, and gateway/routes/allowlist.py trims them to the LLM data-plane surface, so /model/info answers only on the backend and 404s on every replica. All five tests failed at their first read-back in CI while passing against a monolith, where one process serves both planes. The stored row has one answer behind it, so it is read through the shared transport, which routes control-plane paths to the backend. What every gateway must agree on is which models it serves, so the create and delete steps poll /v1/models per replica instead, a route the gateway does serve. read_back_everywhere now rejects a control-plane path outright rather than timing out on it. Two things surfaced behind that. /public/ was missing from the transport's control-plane prefixes, so model_cost_map() was routed to a gateway and 404'd, and the billing steps needed a data-plane wait: a PATCH lands on the backend and each gateway picks it up on its own config reload, measured here at 12-24s, so they now drive calls until the new rate reaches the spend row and let the deadline fail them. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01C1S92J8gSxxKVe1JBzxWBF * test(models): keep polling outcomes immutable and document shared ownership * test: validate opaque stream IDs and hide log-reader credentials * test: isolate auto-router scenarios and clean partial setup --------- Co-authored-by: Claude Opus 5 --- .../model_management_endpoints.py | 70 +++- litellm/router.py | 26 +- tests/e2e/coverage_registry/mgmt.yaml | 2 + .../management/test_model_lifecycle_e2e.py | 365 ++++++++++++++++++ tests/e2e/models.py | 74 +++- tests/e2e/proxy_client.py | 199 +++++++++- tests/e2e/test_proxy_client.py | 57 ++- tests/e2e/transport.py | 1 + .../test_model_management_endpoints.py | 178 ++++++++- .../test_router_model_cost_isolation.py | 192 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 11 files changed, 1130 insertions(+), 39 deletions(-) create mode 100644 tests/e2e/management/test_model_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..742b9d9817f 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -119,6 +119,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, ) @@ -728,6 +729,44 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) +_OWNERSHIP_FIELDS: Final = frozenset( + { + "db_model", + "team_id", + "team_public_model_name", + "access_groups", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + } +) + +_STORED_REQUIRED_FIELDS: Final = frozenset( + name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() +) + +_NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS | frozenset(PTU_MODEL_INFO_FIELDS) + + +def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: + """The keys a patch sends as an explicit null, which update_db_model removes from the + stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored + models require, since clearing one writes a row no reload can rebuild through + LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, + whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob + the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. + """ + if patch is None: + return frozenset() + return frozenset( + field + for field in patch.model_fields_set + if field not in _NULL_CLEAR_IGNORED_FIELDS and getattr(patch, field) is None + ) + + def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -748,25 +787,15 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.model_info: merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) - # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI - # passes through (which today re-sends the OLD pricing on every save) cannot - # silently undo a litellm_params clear via .update(). - # - # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character - # and cache read/write costs) so this path cannot be used to null out privileged - # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are - # mirrored between litellm_params and model_info by Deployment.__init__, so the - # clear propagates to both blobs. - if updated_patch.litellm_params: - for field in updated_patch.litellm_params.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_litellm_params.pop(field, None) - merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.litellm_params): + merged_litellm_params.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.model_info): + merged_model_info.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_litellm_params.pop(field, None) if updated_patch.model_info: - for field in updated_patch.model_info.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_model_info.pop(field, None) - merged_litellm_params.pop(field, None) for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) @@ -816,8 +845,9 @@ async def patch_model( """ PATCH Endpoint for partial model updates. - Only updates the fields specified in the request while preserving other existing values. - Follows proper PATCH semantics by only modifying provided fields. + JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + value is set (identity and ownership keys such as `id` and `team_id` ignore a null). Args: model_id: The ID of the model to update diff --git a/litellm/router.py b/litellm/router.py index 934a4ac86a9..1252d7e7487 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -628,6 +628,15 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat +# namespace with the built-in model catalog. Only a key it created may be evicted, or a +# deployment whose id names a real model would strip that model's pricing and +# capabilities for every other deployment of it. delete_deployment gives a key back once no +# live router still serves that id, so a later catalog refresh that starts serving the name +# is not treated as a deployment's own. +_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys + + class FallbackAwareStreamWrapper(CustomStreamWrapper): """Base for the Router's chat-completion stream wrappers, which are built around the attempt the Router picked first and have to repoint themselves when a fallback takes over.""" @@ -9761,6 +9770,7 @@ class Router: """ idx: Final = len(self.model_list) self.model_list.append(model) + _live_routers.add(self) # mutable-ok: track dynamic routers without extending their lifetimes self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() @@ -9929,7 +9939,12 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. + the entries a refresh rebuilds are the ones a fresh boot would produce. An + entry this function created is replaced rather than merged, so a price cleared + from the deployment does not linger from an earlier registration and keep + billing at the old rate. An entry it did not create is left to merge, because + a deployment id that collides with a catalog model name shares that model's + entry with every other deployment of it. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9946,6 +9961,10 @@ class Router: } if model_id is not None: + if model_id in _DEPLOYMENT_COST_MAP_KEYS: + litellm.model_cost.pop(model_id, None) # mutable-ok: remove cleared prices from the shared entry + elif model_id not in litellm.model_cost: + _DEPLOYMENT_COST_MAP_KEYS.add(model_id) # mutable-ok: retain shared ownership across reloads litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, @@ -10042,6 +10061,11 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) + if not any( + router is not self and id in router.model_id_to_deployment_index_map + for router in tuple(_live_routers) + ): + _DEPLOYMENT_COST_MAP_KEYS.discard(id) # mutable-ok: the last owning router released this key try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index c8d7037d2fd..83f1711a245 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,6 +76,8 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.model.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "model_management_endpoints.py:731", rationale: "A partial PATCH changes only the keys it names; every other stored key reads back byte-for-byte, and the new rate reaches billing"} +- {id: mgmt.model.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "model_management_endpoints.py:731", fail_before_fix: proven, rationale: "An explicit null on PATCH removes the key from the stored row and billing falls back to the cost map; before the fix only mirrored pricing keys could be cleared"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/management/test_model_lifecycle_e2e.py b/tests/e2e/management/test_model_lifecycle_e2e.py new file mode 100644 index 00000000000..99044efeb46 --- /dev/null +++ b/tests/e2e/management/test_model_lifecycle_e2e.py @@ -0,0 +1,365 @@ +"""Live e2e: the lifecycle of a DB-stored deployment through the model management +routes, read back on every gateway replica. + +Each test registers its own gpt-4o-mini mock deployment through /model/new (deleted +on teardown) with non-default pricing, context window, mode, and api_base pinned, then +walks the lifecycle up to the step it proves: the create reads back field for field, +a partial PATCH changes only the key it names, an explicit null on PATCH removes the +key from the stored row (JSON Merge Patch), a call after the price clear is billed at +the cost map's rate rather than the cleared override, and a delete removes the +deployment from /model/info and makes the model name unknown to /chat/completions. + +The stored row is read back from /model/info, a control-plane route with one answer +behind it. What every gateway must agree on is which models it serves, so the create +and delete steps poll /v1/models on every URL in PROXY_REPLICA_URLS through +ProxyClient.read_model_back_everywhere, failing by name on the gateway that never converged. +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + ChatBody, + ChatMessage, + Clear, + LiteLLMParamsBody, + LiteLLMParamsPatch, + ModelInfoBody, + ModelInfoEntry, + ModelInfoResponse, + ModelNewBody, + ModelPatchBody, + ModelsListResponse, + SpendLogRow, +) + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL: Final = "gpt-4o-mini" +PINNED_API_BASE: Final = "https://pinned.example.invalid/v1" +PINNED_MAX_INPUT_TOKENS: Final = 4096 +PINNED_INPUT_RATE: Final = 1e-05 +UPDATED_INPUT_RATE: Final = 2e-05 +PINNED_OUTPUT_RATE: Final = 3e-05 + +# A PATCH lands on the control plane, and each gateway picks it up on its own config +# reload, so the first call after the write can still be billed at the old rate. There +# is no price on the gateway's data-plane surface to poll, so the billing steps drive +# calls until the new rate shows up in the spend row and let the deadline be what fails. +BILLING_CONVERGENCE_TIMEOUT: Final = 90.0 +BILLING_CONVERGENCE_INTERVAL: Final = 5.0 + + +@dataclass(frozen=True, slots=True) +class Registered: + model_name: str + model_id: str + + +class _ErrorDetail(BaseModel): + message: str + + +class _ErrorEnvelope(BaseModel): + error: _ErrorDetail + + +def _register(client: ManagementClient, resources: ResourceManager) -> Registered: + """Register a mock gpt-4o-mini deployment with every field under test pinned to a + non-default value, deleted on teardown. max_input_tokens is pinned in + litellm_params only: a value in model_info is copied into the shared cost-map + entry for the backend model, which would leak into every other gpt-4o-mini + deployment on the proxy.""" + model_name: Final = f"e2e-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=BACKEND_MODEL, + mock_response="ok", + api_base=PINNED_API_BASE, + input_cost_per_token=PINNED_INPUT_RATE, + output_cost_per_token=PINNED_OUTPUT_RATE, + max_input_tokens=PINNED_MAX_INPUT_TOKENS, + ), + model_info=ModelInfoBody(mode="chat"), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return Registered(model_name=model_name, model_id=model_id) + + +def _entry(body: ModelInfoResponse, model_name: str) -> ModelInfoEntry | None: + return next((entry for entry in body.data if entry.model_name == model_name), None) + + +def _stored_entry( + client: ManagementClient, + model_name: str, + *, + converged: Callable[[ModelInfoEntry], bool], +) -> ModelInfoEntry: + """The stored /model/info row for `model_name`, once it satisfies `converged`. + + /model/info is a control-plane route: the gateways named in PROXY_REPLICA_URLS + serve the LLM surface only, so the stored row has one answer, not one per + gateway. What every gateway must agree on is which models it serves, and + `_assert_served_everywhere` / `_assert_absent_everywhere` poll /v1/models for + that.""" + + def has_converged(body: ModelInfoResponse) -> bool: + entry: Final = _entry(body, model_name) + return entry is not None and converged(entry) + + body: Final = client.proxy.read_model_back("/model/info", ModelInfoResponse, predicate=has_converged) + entry: Final = _entry(body, model_name) + assert entry is not None, f"/model/info stopped listing {model_name!r} between the poll and the read" + return entry + + +def _serves(body: ModelsListResponse, model_name: str) -> bool: + return any(entry.id == model_name for entry in body.data) + + +def _assert_served_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: _serves(body, model_name) + ) + + +def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: not _serves(body, model_name) + ) + + +def _assert_untouched_keys_as_created(entry: ModelInfoEntry, replica: str) -> None: + """The keys no later step names read back byte-for-byte as /model/new wrote them.""" + params: Final = entry.litellm_params + assert params.model == BACKEND_MODEL, f"{replica}: litellm_params.model {params.model!r} != {BACKEND_MODEL!r}" + assert params.api_base == PINNED_API_BASE, f"{replica}: api_base {params.api_base!r} != {PINNED_API_BASE!r}" + assert params.output_cost_per_token == PINNED_OUTPUT_RATE, ( + f"{replica}: output_cost_per_token {params.output_cost_per_token} != {PINNED_OUTPUT_RATE}" + ) + assert entry.model_info.mode == "chat", f"{replica}: model_info.mode {entry.model_info.mode!r} != 'chat'" + + +def _approx_equal(actual: float, expected: float) -> bool: + return math.isclose(actual, expected, rel_tol=1e-2, abs_tol=1e-9) + + +def _priced(rows: list[SpendLogRow]) -> bool: + return any(row.metadata and row.metadata.cost_breakdown and row.metadata.cost_breakdown.input_cost for row in rows) + + +def _billed_input_cost(client: ManagementClient, model_name: str, key: str) -> tuple[int, float]: + """Drive one chat completion through `model_name` and return the prompt tokens and + input cost its spend row recorded, so a test can assert the rate the gateway actually + billed rather than only the rate it stored.""" + chat: Final = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], + max_tokens=16, + ), + ) + ) + assert chat.id is not None, f"chat completion carried no id to find its spend row by: {chat}" + + rows: Final = client.proxy.poll_logs_for_request_id(chat.id, predicate=_priced) + row: Final = next((row for row in rows if row.request_id == chat.id), None) + assert row is not None and row.metadata and row.metadata.cost_breakdown, ( + f"no priced spend row for request {chat.id} before the deadline: {rows}" + ) + prompt_tokens: Final = row.prompt_tokens or 0 + input_cost: Final = row.metadata.cost_breakdown.input_cost + assert prompt_tokens > 0 and input_cost is not None, f"spend row logged no prompt tokens or input cost: {row}" + return prompt_tokens, input_cost + + +def _await_billed_input_cost( + client: ManagementClient, model_name: str, key: str, *, expected_rate: float +) -> tuple[int, float]: + """Drive calls through `model_name` until one is billed at `expected_rate`, and + return the prompt tokens and input cost of the last spend row either way. + + Only the deadline ends the wait unsatisfied: a rate that never reaches the gateway + comes back as the stale cost for the caller to assert on, so the rate the caller + expects is still what decides the test.""" + deadline: Final = time.monotonic() + BILLING_CONVERGENCE_TIMEOUT + while True: + prompt_tokens, input_cost = _billed_input_cost(client, model_name, key) + if _approx_equal(input_cost, prompt_tokens * expected_rate) or time.monotonic() >= deadline: + return prompt_tokens, input_cost + time.sleep(BILLING_CONVERGENCE_INTERVAL) + + +class TestModelLifecycle: + @pytest.mark.covers("mgmt.model.add.persists") + def test_create_reads_back_every_field_and_serves_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + entry = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( + f"{stored}: input_cost_per_token {entry.litellm_params.input_cost_per_token} != {PINNED_INPUT_RATE}" + ) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.id == registered.model_id, ( + f"{stored}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" + ) + + _assert_served_everywhere(client, registered.model_name) + + @pytest.mark.covers("mgmt.model.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_key( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(input_cost_per_token=UPDATED_INPUT_RATE)), + ) + assert stored.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"PATCH response stores input_cost_per_token {stored.litellm_params.input_cost_per_token}, " + f"sent {UPDATED_INPUT_RATE}" + ) + + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} " + f"did not mirror the updated {UPDATED_INPUT_RATE}" + ) + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=UPDATED_INPUT_RATE + ) + assert _approx_equal(input_cost, prompt_tokens * UPDATED_INPUT_RATE), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * updated rate {UPDATED_INPUT_RATE} " + f"= {prompt_tokens * UPDATED_INPUT_RATE}; the partial update did not reach billing" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_explicit_null_removes_the_key_from_the_stored_row( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + stored_params = stored.litellm_params.model_fields_set + assert "max_input_tokens" not in stored_params, ( + f"stored litellm_params still carries max_input_tokens " + f"{stored.litellm_params.max_input_tokens} after an explicit null" + ) + assert "input_cost_per_token" not in stored_params, ( + f"stored litellm_params still carries input_cost_per_token " + f"{stored.litellm_params.input_cost_per_token} after an explicit null" + ) + assert "max_input_tokens" not in stored.model_info.model_fields_set, ( + f"stored model_info carries max_input_tokens {stored.model_info.max_input_tokens} after the clear" + ) + assert "input_cost_per_token" not in stored.model_info.model_fields_set, ( + f"stored model_info still mirrors input_cost_per_token {stored.model_info.input_cost_per_token}" + ) + + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + served = entry.litellm_params.model_fields_set + assert "max_input_tokens" not in served, ( + f"{stored}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" + ) + assert "input_cost_per_token" not in served, ( + f"{stored}: litellm_params still serves input_cost_per_token {entry.litellm_params.input_cost_per_token}" + ) + assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} is not the " + f"cost map's {cost_map_input_rate}; the cleared override {PINNED_INPUT_RATE} still resolves" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_cleared_price_is_billed_at_the_cost_map_rate( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + _ = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "input_cost_per_token" not in entry.litellm_params.model_fields_set, + ) + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=cost_map_input_rate + ) + + assert _approx_equal(input_cost, prompt_tokens * cost_map_input_rate), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * cost map rate {cost_map_input_rate} " + f"= {prompt_tokens * cost_map_input_rate}" + ) + assert not _approx_equal(input_cost, prompt_tokens * PINNED_INPUT_RATE), ( + f"input_cost {input_cost} is still billed at the cleared override {PINNED_INPUT_RATE}" + ) + + @pytest.mark.covers("mgmt.model.delete.persists") + def test_delete_removes_the_deployment_everywhere( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + + client.delete_model_strict(registered.model_id) + + _assert_absent_everywhere(client, registered.model_name) + refused = client.chat_status(scoped_key, registered.model_name, "hi this is a test") + assert refused.status_code == 400, ( + f"chat against the deleted model must be rejected 400, got {refused.status_code}: {refused.body[:300]}" + ) + envelope = _ErrorEnvelope.model_validate_json(refused.body) + assert envelope.error.message, f"400 body must carry an error message: {refused.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..8db37bd25a5 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,6 +655,11 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + class GuardrailEntityMatch(BaseModel): entity_type: str score: float @@ -672,6 +677,7 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): + cost_breakdown: CostBreakdown | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None @@ -815,15 +821,42 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token +class DeploymentParams(CustomPricing): + """The litellm_params half of a /model/info row: the stored deployment as written, + credentials scrubbed. Unlike model_info it is never back-filled from the cost map, + so a key the store dropped is absent here (check `model_fields_set`).""" + + model: str | None = None + api_base: str | None = None + max_input_tokens: int | None = None + + +class DeploymentModelInfo(CustomPricing): + id: str | None = None + max_input_tokens: int | None = None + + class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for - it - the override merged over the cost-map defaults.""" + it - the override merged over the cost-map defaults, so a key cleared from the + stored blob reads as the cost-map default here.""" model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: CustomPricing = CustomPricing() - model_info: CustomPricing = CustomPricing() + litellm_params: DeploymentParams = DeploymentParams() + model_info: DeploymentModelInfo = DeploymentModelInfo() + + +class StoredDeployment(BaseModel): + """PATCH /model/{model_id}/update answers with the row as stored: both blobs raw, + nothing back-filled, so a cleared key is absent from `model_fields_set` of the + blob it was cleared from.""" + + model_config = ConfigDict(protected_namespaces=()) + model_name: str + litellm_params: DeploymentParams + model_info: DeploymentModelInfo class ModelInfoResponse(BaseModel): @@ -924,9 +957,10 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None + max_input_tokens: int | None = None -ModelMode = Literal["batch", "realtime", "image_generation"] +ModelMode = Literal["chat", "batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): @@ -936,6 +970,7 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None + max_input_tokens: int | None = None access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None @@ -964,6 +999,37 @@ class ModelUpdateBody(BaseModel): model_info: ModelInfoBody +class Clear(BaseModel): + """Serializes to JSON null. The transport dumps every body with exclude_none, so a + field set to this is how a patch carries the explicit null that removes a stored key.""" + + @model_serializer + def _as_null(self) -> None: + return None + + +class LiteLLMParamsPatch(BaseModel): + api_base: str | Clear | None = None + max_input_tokens: int | Clear | None = None + input_cost_per_token: float | Clear | None = None + output_cost_per_token: float | Clear | None = None + + +class ModelInfoPatch(BaseModel): + mode: ModelMode | Clear | None = None + max_input_tokens: int | Clear | None = None + + +class ModelPatchBody(BaseModel): + """PATCH /model/{model_id}/update body, JSON Merge Patch over the stored deployment: + a field left None is dropped from the body and unchanged, a field set to `Clear()` + is sent as null and removed, a field with a value is set.""" + + model_config = ConfigDict(protected_namespaces=()) + litellm_params: LiteLLMParamsPatch | None = None + model_info: ModelInfoPatch | None = None + + class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1bac5116a9d..fa1b06fe7ed 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,10 +10,10 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass -from functools import reduce from datetime import datetime +from functools import reduce from types import MappingProxyType from typing import Final @@ -62,6 +62,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -72,6 +73,7 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + StoredDeployment, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -132,6 +134,103 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +type BodyReader[R: BaseModel] = Callable[[float], Result[R]] + + +@dataclass(frozen=True, slots=True) +class BodyNotConverged[R: BaseModel]: + """The deadline passed without a read the predicate accepted; `last_result` is the + final read, so the caller can tell a body that never matched from a read that + failed.""" + + last_result: Result[R] | None + + +@dataclass(frozen=True, slots=True) +class BodyConverged[R: BaseModel]: + """Every replica answered a body the predicate accepted; `bodies` is the last read + per replica.""" + + bodies: Mapping[str, R] + + +@dataclass(frozen=True, slots=True) +class BodyNeverConvergedOn[R: BaseModel]: + """`BodyNotConverged` labeled with the replica whose reads never satisfied the predicate.""" + + replica: str + last_result: Result[R] | None + + +def await_body_converged[R: BaseModel]( + read: BodyReader[R], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Success[R] | BodyNotConverged[R]: + """Poll `read` until it answers a body `predicate` accepts, or `timeout` passes. + + Each read's request timeout is clamped to the remaining budget, and the sleep + between reads to the time left, so the last read before the deadline is never + skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + + def reads() -> Iterator[Result[R]]: + while (remaining := deadline - now()) > 0: + yield read(min(request_timeout, remaining)) + sleep(min(interval, max(deadline - now(), 0.0))) + + def attempts() -> Iterator[Success[R] | BodyNotConverged[R]]: + for result in reads(): + if isinstance(result, Success) and predicate(result.data): + yield result + return + yield BodyNotConverged(last_result=result) + + initial: Final[Success[R] | BodyNotConverged[R]] = BodyNotConverged(last_result=None) + return reduce(lambda _previous, result: result, attempts(), initial) + + +def await_body_converged_everywhere[R: BaseModel]( + readers: Mapping[str, BodyReader[R]], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + """`await_body_converged` against every replica in turn, each with the full budget, so a + write counts as landed only once every replica serves it.""" + def read_replica( + outcome: BodyConverged[R] | BodyNeverConvergedOn[R], + item: tuple[str, BodyReader[R]], + ) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + if isinstance(outcome, BodyNeverConvergedOn): + return outcome + replica, read = item + match await_body_converged( + read, + predicate=predicate, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ): + case Success(data=data): + return BodyConverged(bodies=MappingProxyType({**outcome.bodies, replica: data})) + case BodyNotConverged(last_result=last_result): + return BodyNeverConvergedOn(replica=replica, last_result=last_result) + initial: Final[BodyConverged[R] | BodyNeverConvergedOn[R]] = BodyConverged(bodies=MappingProxyType({})) + return reduce(read_replica, readers.items(), initial) + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -658,6 +757,102 @@ class ProxyClient: ) ) + def patch_model(self, model_id: str, body: ModelPatchBody) -> StoredDeployment: + """JSON Merge Patch the deployment `model_id` via PATCH /model/{model_id}/update: + a field the body omits is unchanged, one sent as null is removed from the stored + row, one sent with a value is set. See ModelPatchBody for how a null is sent. + Returns the row as stored after the write.""" + return unwrap( + self.transport.patch( + f"/model/{model_id}/update", + headers=self.transport.master, + json=body, + response_type=StoredDeployment, + ) + ) + + def read_model_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, predicate: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica until each answers a body `predicate` accepts, + polling to poll_timeout, and return the last body per replica. + + Fails naming the replica that never converged, so a write that reached one + gateway but not the others is caught instead of passing on whichever gateway + the balancer answered from. Falls back to the single proxy address when no + replica list is configured. + + `path` must be a data-plane route. The replicas are gateways, which serve only + the LLM surface, so a control-plane path answers on exactly one service and + 404s on every replica in a split deployment: asking each replica for one is + never the question the caller means. Read those through `self.transport` + instead, which routes them to the control plane.""" + if is_control_plane_path(path): + raise AssertionError( + f"read_model_back_everywhere({path!r}) asks every data-plane replica for a control-plane route. " + "The replicas are gateways and do not serve it; poll a data-plane path such as /v1/models " + "here, and read the control plane through the shared transport." + ) + readers: Final = { + url: self._body_reader(transport, path, response_type) + for url, transport in self._read_back_replicas().items() + } + outcome: Final = await_body_converged_everywhere( + readers, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies + case BodyNeverConvergedOn(replica=replica, last_result=last_result): + raise AssertionError( + f"GET {path} on {replica} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def read_model_back[R: BaseModel](self, path: str, response_type: type[R], *, predicate: Callable[[R], bool]) -> R: + """GET `path` through the shared transport until the body satisfies `predicate`, + polling to poll_timeout, and return that body. + + The counterpart to `read_model_back_everywhere` for a control-plane route such as + /model/info: the stored row lives in one database behind one control plane, so + there is a single answer to converge on rather than one per gateway.""" + outcome: Final = await_body_converged_everywhere( + {CONTROL_PLANE_BASE_URL: self._body_reader(self.transport, path, response_type)}, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies[CONTROL_PLANE_BASE_URL] + case BodyNeverConvergedOn(last_result=last_result): + raise AssertionError( + f"GET {path} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def _read_back_replicas(self) -> Mapping[str, Transport]: + return self.replicas or MappingProxyType({CONTROL_PLANE_BASE_URL: self.transport}) + + @staticmethod + def _body_reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> BodyReader[R]: + return lambda timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=timeout, + ) + def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 3b84a47e3cc..cbf7f5648d4 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -20,8 +20,12 @@ from typing import Final, cast import pytest from e2e_config import parse_replica_urls from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from models import KeyInfo, KeyInfoResponse, ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( + BodyReader, + BodyConverged, + BodyNeverConvergedOn, + await_body_converged_everywhere, ConvergeOutcome, Converged, EverywhereConverged, @@ -274,3 +278,54 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +def _info(*model_names: str) -> Success[ModelInfoResponse]: + entries: Final = [ModelInfoEntry(model_name=model_name) for model_name in model_names] + return Success(status_code=200, data=ModelInfoResponse(data=entries)) + + +def _reader(results: Iterable[Success[ModelInfoResponse]]) -> BodyReader[ModelInfoResponse]: + it: Final = iter(results) + return lambda _timeout: next(it) + + +def _lists_model(body: ModelInfoResponse) -> bool: + return any(entry.model_name == MODEL for entry in body.data) + + +def _read_back( + readers: Mapping[str, BodyReader[ModelInfoResponse]], +) -> tuple[BodyConverged[ModelInfoResponse] | BodyNeverConvergedOn[ModelInfoResponse], FakeClock]: + clock: Final = FakeClock() + outcome: Final = await_body_converged_everywhere( + readers, + predicate=_lists_model, + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + return outcome, clock + + +class TestAwaitBodyConvergedEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_body(self) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(chain(repeat(_info(), 2), repeat(_info(MODEL)))), + } + outcome, clock = _read_back(readers) + assert outcome == BodyConverged(bodies={"gateway-1": _info(MODEL).data, "gateway-2": _info(MODEL).data}) + assert clock.elapsed == 2 * INTERVAL + + @pytest.mark.parametrize("lagging", ["gateway-1", "gateway-2"]) + def test_fails_naming_the_replica_that_never_converges(self, lagging: str) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(repeat(_info(MODEL))), + } | {lagging: _reader(repeat(_info()))} + outcome, clock = _read_back(readers) + assert outcome == BodyNeverConvergedOn(replica=lagging, last_result=_info()) + assert clock.elapsed >= TIMEOUT diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..804e073a4a0 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,6 +306,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/config", "/guardrails", "/openapi.json", + "/public/", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index c02f886fc31..1376727e296 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3115,9 +3115,6 @@ class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored between the two by Deployment.__init__). - - Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) - cannot be cleared via this path. """ def test_clear_input_cost_removes_from_both_blobs(self): @@ -3193,10 +3190,10 @@ class TestUpdateDBModelClearPricing: assert params["input_cost_per_token"] == 0.000001 assert params["output_cost_per_token"] == 0.000007 - def test_null_on_non_pricing_field_does_not_clear(self): - """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. - Privileged or unrelated model_info fields (e.g. team_id) must be unaffected - by the null-clearing path so a team admin can't ungate a team-scoped model. + def test_null_on_one_field_leaves_other_fields_alone(self): + """A null clears only the key it names: pricing the patch never mentions and + the ownership key team_id stay put, so a team admin can't ungate a + team-scoped model through the clear path. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, @@ -3217,8 +3214,6 @@ class TestUpdateDBModelClearPricing: model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), ) - # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id - # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, updated_patch=updateDeployment( @@ -3394,6 +3389,171 @@ class TestUpdateDBModelClearPricing: assert info["cache_creation_input_token_cost"] == 0.000003 +_PROTECTED_MODEL_INFO_VALUES = { + "team_id": "team-keep-me", + "team_public_model_name": "team-facing-name", + "access_groups": ["group-a"], + "created_at": "2026-01-01T00:00:00+00:00", + "created_by": "creator", + "updated_at": "2026-01-02T00:00:00+00:00", + "updated_by": "updater", + "blocked": True, +} + + +def _build_db_model_with_pinned_model_info(): + """Deployment whose stored blobs pin non-pricing keys an earlier save wrote, next to a + pricing override, so a clear can be checked key by key.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="pinned-gpt-4o-mini", + litellm_params=LiteLLM_Params( + model="gpt-4o-mini", input_cost_per_token=0.000001, max_input_tokens=4096 + ), + model_info=ModelInfo( + id="dep-pinned-0", + max_input_tokens=4096, + mode="chat", + supports_vision=True, + **_PROTECTED_MODEL_INFO_VALUES, + ), + ) + + +class TestUpdateDBModelNullClearsAnyKey: + """JSON Merge Patch on PATCH /model/{id}/update: a key sent as null is removed from the + stored blob it was sent in, whatever the key, except the identity and ownership keys, + whose nulls are ignored.""" + + def test_model_info_nulls_remove_pinned_non_pricing_keys(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"max_input_tokens": None, "mode": None}} + ), + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info + assert "mode" not in info + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + + def test_litellm_params_null_removes_pinned_non_pricing_key(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"litellm_params": {"max_input_tokens": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in params + assert params["model"] == "gpt-4o-mini" + assert params["input_cost_per_token"] == 0.000001 + assert info["max_input_tokens"] == 4096 + + def test_omitted_key_is_untouched_by_a_null_elsewhere(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"mode": None, "supports_vision": False}} + ), + ) + + info = json.loads(result["model_info"]) + assert "mode" not in info + assert info["supports_vision"] is False + assert info["max_input_tokens"] == 4096 + + @pytest.mark.parametrize("field", sorted(_PROTECTED_MODEL_INFO_VALUES)) + def test_null_on_protected_key_is_ignored(self, field): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate({"model_info": {field: None}}), + ) + + info = json.loads(result["model_info"]) + assert info[field] == _PROTECTED_MODEL_INFO_VALUES[field] + assert info["max_input_tokens"] == 4096 + + def test_echoing_the_read_back_blob_preserves_every_stored_key(self): + """The Admin UI edit form submits the whole /model/info row back, and that read reports + every key the deployment never stored as an explicit null. Those nulls have to stay + no-ops: a write drops None before storing, so a null in the echoed blob always names a + key the stored row does not carry. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + db_model = _build_db_model_with_pinned_model_info() + echoed = { + "id": "dep-pinned-0", + "max_input_tokens": 4096, + "mode": "chat", + "supports_vision": True, + "input_cost_per_token": 0.000001, + "team_id": "team-keep-me", + "base_model": None, + "tier": None, + "max_output_tokens": None, + "supports_function_calling": None, + "cache_read_input_token_cost": None, + } + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment.model_validate({"model_info": echoed}), + ) + + info = json.loads(result["model_info"]) + assert info["max_input_tokens"] == 4096 + assert info["mode"] == "chat" + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + assert info["team_id"] == "team-keep-me" + for never_stored in ("base_model", "tier", "max_output_tokens", "supports_function_calling"): + assert never_stored not in info + + def test_null_on_pricing_key_still_clears_both_blobs(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"input_cost_per_token": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + assert params["max_input_tokens"] == 4096 + assert info["max_input_tokens"] == 4096 + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 30b265905f3..d22ec60e61a 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,6 +220,198 @@ def test_should_store_full_pricing_under_deployment_model_id(): assert entry["output_cost_per_token"] == 0.0 +def test_should_drop_a_price_the_deployment_no_longer_carries(): + """Re-registering a deployment must replace its model_id entry, not merge onto it. + + A merge left the old rate in the cost map after an operator cleared the override, so + the deployment kept billing at a price its config no longer had. + """ + backend_model = "vertex_ai/gemini-2.5-flash" + model_id = "deployment-cleared-price" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"input_cost_per_token": 0.005, "output_cost_per_token": 0.01}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"mode": "chat"}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + + entry = litellm.model_cost[model_id] + assert entry.get("input_cost_per_token") != 0.005, ( + "the cleared override survived re-registration, so the deployment still bills at it" + ) + assert entry.get("output_cost_per_token") != 0.01 + finally: + _restore_model_cost_entries(original) + + +def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it(): + """Deployments are keyed into the same cost map as the built-in catalog, so a deployment + whose id happens to name a real model must not evict that model's entry. + + Stripping it would take the pricing and capability flags every other deployment of that + model reads, process-wide, until the next price-map reload. Registering twice, because + the first registration is what would mark the entry as this deployment's own. + """ + colliding_id = "gpt-4o" + original = {colliding_id: litellm.model_cost.get(colliding_id)} + builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] + + try: + for _ in range(2): + Router._register_deployment_in_model_cost( + model_id=colliding_id, + model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + entry = litellm.model_cost[colliding_id] + assert entry["max_tokens"] == builtin_max_tokens, ( + "registering a deployment under a catalog model's name wiped that model's context window" + ) + assert entry["litellm_provider"] == "openai" + assert entry["supports_vision"] is True + finally: + _restore_model_cost_entries(original) + + +def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): + """A deployment may carry `litellm_provider` in its own model_info, which must not be + read as "this is a catalog entry" and stop the stale price from being dropped.""" + model_id = "deployment-provider-tagged" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( + "a deployment that declares its provider kept billing at the price it no longer carries" + ) + finally: + _restore_model_cost_entries(original) + + +def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): + """Deleting a deployment releases its claim on the shared cost-map key. + + Held forever, a later catalog refresh that starts publishing a model under that same + name would be treated as the deleted deployment's own entry and evicted. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-to-delete" + original = {model_id: litellm.model_cost.get(model_id)} + router = Router( + model_list=[ + { + "model_name": "to-delete", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + ] + ) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert router.delete_deployment(id=model_id) is not None + + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS, ( + "a deleted deployment kept its claim on the shared cost-map key" + ) + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): + """Two live routers can serve the same deployment id, and the claim is process-wide. + + Releasing it when only one of them drops the deployment would put the survivor back on + merging, so the price it just cleared would keep billing. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-served-twice" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "served-twice", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + first = Router(model_list=[entry]) + second = Router(model_list=[entry]) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert first.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while another router still served the deployment" + ) + + assert second.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_a_dynamically_built_router_serves_it(): + """A router built with no model_list still serves whatever add_deployment gives it, so it + counts when deciding whether the shared cost-map claim can be released.""" + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-added-dynamically" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "added-dynamically", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + configured = Router(model_list=[entry]) + dynamic = Router() + dynamic.add_deployment(deployment=Deployment(**entry)) + + try: + assert configured.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while a dynamically built router still served the deployment" + ) + + assert dynamic.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..93aaf3ca58c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9064,8 +9064,9 @@ export interface paths { * Patch Model * @description PATCH Endpoint for partial model updates. * - * Only updates the fields specified in the request while preserving other existing values. - * Follows proper PATCH semantics by only modifying provided fields. + * JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + * body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + * value is set (identity and ownership keys such as `id` and `team_id` ignore a null). * * Args: * model_id: The ID of the model to update From 25a6c2a2749ba802514fb3a5736b540c21a1e630 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:17:59 +0000 Subject: [PATCH 088/157] fix(cli): skip remote model cost map fetch in lite CLI processes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 14 ++++- .../test_get_model_cost_map.py | 25 +++++++++ .../proxy/client/cli/test_global_options.py | 56 ++++++++++++++++++- 3 files changed, 91 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index cdc4810ff04..7d3b67e58c8 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -2,6 +2,7 @@ Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True. +The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching. ``` export LITELLM_LOCAL_MODEL_COST_MAP=True @@ -13,11 +14,13 @@ import hashlib import json import os import random +import sys import time from collections.abc import Awaitable, Callable from dataclasses import dataclass, replace from datetime import datetime, timezone from importlib.resources import files +from pathlib import Path from typing import Final, Protocol import httpx @@ -33,6 +36,12 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) + + +def _is_cli_process() -> bool: + return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. @@ -531,7 +540,8 @@ def get_model_cost_map( """ Public entry point — returns the model cost map dict. - 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. + 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` / + ``litellm-proxy`` CLI process, uses the local backup only. 2. Otherwise fetches from ``url``, retrying transient HTTP errors (429/5xx/transport) with Retry-After-aware backoff, validates integrity, and falls back to the local backup on any failure. @@ -543,7 +553,7 @@ def get_model_cost_map( _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 0495440c51c..d092b387259 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import sys import pytest @@ -711,3 +712,27 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej assert source["etag"] is None assert source["source_revision"] == _bundled_blob_id() assert source["source_revision"] != git_blob_id(shrunk_body) + + +@pytest.mark.parametrize( + ("argv0", "request_count"), + [ + ("/some/venv/bin/lite", 0), + ("/some/venv/bin/lite.exe", 0), + ("/some/venv/bin/python", 1), + ], +) +def test_boot_load_skips_remote_fetch_for_cli_processes(monkeypatch, argv0, request_count): + monkeypatch.setattr(sys, "argv", [argv0, "--version"]) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + + cost_map = get_model_cost_map(url=_URL, client=client) + + assert calls["count"] == request_count + assert cost_map + source = get_model_cost_map_source_info() + if request_count == 0: + assert source["source"] == "local" + else: + assert source["source"] == "remote" diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 0dd388919a5..02b1cdf4296 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,14 +1,16 @@ # stdlib imports import json import os +import shutil +import subprocess +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest.mock import Mock, patch import pytest from click.testing import CliRunner - - import litellm.proxy.client.cli from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli @@ -35,6 +37,56 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output +def test_lite_version_does_not_fetch_model_cost_map(): + lite_path = shutil.which("lite") + if lite_path is None: + pytest.skip("lite executable is unavailable") + + requests = [] + + class _CostMapHandler(BaseHTTPRequestHandler): + def do_GET(self): + requests.append(self.path) + body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format, *args): + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + source_root = str(Path(__file__).resolve().parents[5]) + env["PYTHONPATH"] = os.pathsep.join(filter(None, (source_root, env.get("PYTHONPATH")))) + env.update( + { + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": "http://127.0.0.1:9", + } + ) + result = subprocess.run( + [lite_path, "--version"], + capture_output=True, + text=True, + timeout=120, + env=env, + ) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=10) + + assert result.returncode == 0 + assert "LiteLLM Proxy CLI Version" in result.stdout + assert requests == [] + + def test_cli_source_is_ascii_only(): """Non-ASCII output (emoji, box-drawing chars) raises UnicodeEncodeError on legacy Windows consoles (cp1252), so the whole CLI package must stay ASCII-only.""" From ffae447649e173214a2081b4890410a11dfdf085 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:30:50 +0000 Subject: [PATCH 089/157] test(cli): type cost map bypass regression handlers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/test_get_model_cost_map.py | 4 +++- tests/test_litellm/proxy/client/cli/test_global_options.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index d092b387259..6ab3f9a21af 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -722,7 +722,9 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej ("/some/venv/bin/python", 1), ], ) -def test_boot_load_skips_remote_fetch_for_cli_processes(monkeypatch, argv0, request_count): +def test_boot_load_skips_remote_fetch_for_cli_processes( + monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int +) -> None: monkeypatch.setattr(sys, "argv", [argv0, "--version"]) monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 02b1cdf4296..6de56573ffb 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -42,10 +42,10 @@ def test_lite_version_does_not_fetch_model_cost_map(): if lite_path is None: pytest.skip("lite executable is unavailable") - requests = [] + requests: list[str] = [] class _CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self): + def do_GET(self) -> None: requests.append(self.path) body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' self.send_response(200) @@ -54,7 +54,7 @@ def test_lite_version_does_not_fetch_model_cost_map(): self.end_headers() self.wfile.write(body) - def log_message(self, format, *args): + def log_message(self, format: str, *args: object) -> None: return server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) From ef3a3c16ae02bc6d83e14b09c437ef36081c0498 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:32:31 -0700 Subject: [PATCH 090/157] feat(guardrails): map each guardrail scan id to its guardrail, stage and provider (#40327) * feat(guardrails): map each guardrail scan id to its guardrail, stage and provider Adds the x-litellm-guardrail-scan-metadata response header, a JSON list of {guardrail, stage, provider, scan_id} entries, next to the existing comma-separated x-litellm-guardrail-scan-id header. Prisma AIRS records the execution stage for every scan and OpenAI Moderation now records its moderation id too. The new metadata key is internal: client-supplied values are stripped and it is exposed through the UI CORS allow list. Resolves LIT-6018 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(guardrails): cap the scan metadata response header at a configurable length Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(guardrails): hardcode the scan metadata header cap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 + litellm/proxy/common_utils/callback_utils.py | 62 ++++++++- .../guardrail_hooks/openai/moderations.py | 10 +- .../panw_prisma_airs/panw_prisma_airs.py | 34 +++-- litellm/proxy/litellm_pre_call_utils.py | 2 + .../proxy/common_utils/test_callback_utils.py | 122 +++++++++++++++--- .../openai/test_moderations.py | 32 +++++ .../guardrail_hooks/test_panw_prisma_airs.py | 32 ++++- 8 files changed, 265 insertions(+), 31 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f2384c8c6c..108a914e9c1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -143,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) +MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048 DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000 @@ -197,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-adaptive-router-model", "x-litellm-applied-guardrails", "x-litellm-guardrail-scan-id", + "x-litellm-guardrail-scan-metadata", "x-litellm-cache-key", ] diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 770963a1f24..561a53409f4 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,10 +1,12 @@ import copy +import json import os from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass +from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias -from typing_extensions import assert_never +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm import get_secret @@ -12,6 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -28,6 +31,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -52,6 +56,15 @@ reset_color_code: Final = "\033[0m" TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY: Final = "_pillar_response_headers_trusted" GUARDRAIL_SCAN_IDS_METADATA_KEY: Final = "guardrail_scan_ids" +GUARDRAIL_SCAN_METADATA_METADATA_KEY: Final = "guardrail_scan_metadata" + + +class GuardrailScanMetadata(TypedDict): + guardrail: ReadOnly[str | None] + stage: ReadOnly[str] + provider: ReadOnly[str] + scan_id: ReadOnly[str] + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -450,6 +463,16 @@ def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, return headers +def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: int) -> str | None: + """Compact JSON list of scan metadata entries, dropping trailing entries so the header fits in max_length.""" + encoded: Final = tuple(json.dumps(entry, separators=(",", ":")) for entry in entries) + lengths: Final = tuple(accumulate(len(item) + 1 for item in encoded)) + kept: Final = sum(1 for length in lengths if length + 1 <= max_length) + if kept == 0: + return None + return f"[{','.join(encoded[:kept])}]" + + def get_logging_caching_headers(request_data: dict) -> dict | None: _metadata: Final[dict] = {} metadata_bucket: Final = request_data.get("metadata") @@ -468,6 +491,15 @@ def get_logging_caching_headers(request_data: dict) -> dict | None: if scan_ids: headers["x-litellm-guardrail-scan-id"] = ",".join(scan_ids) + scan_metadata: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + scan_metadata_header: Final = ( + _serialize_scan_metadata_header(scan_metadata, max_length=MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH) + if isinstance(scan_metadata, (list, tuple)) + else None + ) + if scan_metadata_header: + headers["x-litellm-guardrail-scan-metadata"] = scan_metadata_header + if "applied_policies" in _metadata: headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"]) @@ -501,6 +533,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( "applied_policies", "applied_guardrails", GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, "policy_sources", "guardrails", "guardrail_config", @@ -565,21 +598,40 @@ def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_nam _metadata["applied_guardrails"] = [guardrail_name] -def add_guardrail_scan_id(request_data: dict, scan_id: str | None) -> None: +def add_guardrail_scan_id( + request_data: dict[str, object], + scan_id: str | None, + *, + guardrail_name: str | None, + provider: str, + stage: GuardrailEventHooks, +) -> None: """ - Record a provider scan id so it can be surfaced to the caller. + Record a provider scan id, keyed to the guardrail execution that produced it, so it can be surfaced to the caller. Guardrails only return scan details to the client when they block, so allowed requests carry no - audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header. + audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header, and the + (guardrail, stage, provider, scan_id) entries become the x-litellm-guardrail-scan-metadata header. """ if not scan_id: return _, _metadata = get_or_create_metadata_bucket(request_data) existing: Final = _metadata.get(GUARDRAIL_SCAN_IDS_METADATA_KEY) - scan_ids: Final = tuple(existing) if isinstance(existing, (list, tuple)) else () + scan_ids: Final[tuple[object, ...]] = tuple(existing) if isinstance(existing, (list, tuple)) else () if scan_id not in scan_ids: _metadata[GUARDRAIL_SCAN_IDS_METADATA_KEY] = (*scan_ids, scan_id) + entry: Final[GuardrailScanMetadata] = { + "guardrail": guardrail_name, + "stage": stage.value, + "provider": provider, + "scan_id": scan_id, + } + existing_entries: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + entries: Final[tuple[object, ...]] = tuple(existing_entries) if isinstance(existing_entries, (list, tuple)) else () + if entry not in entries: + _metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry) + def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 10683550f85..c22d35509c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -17,7 +17,8 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id +from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from litellm.types.utils import ( GenericGuardrailAPIInputs, GuardrailStatus, @@ -218,6 +219,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): metadata: Final = request_data.get("metadata") or {} request_data["metadata"] = metadata metadata["_openai_moderation_response"] = moderation_response.model_dump() + add_guardrail_scan_id( + request_data=request_data, + scan_id=moderation_response.id, + guardrail_name=self.guardrail_name, + provider=SupportedGuardrailIntegrations.OPENAI_MODERATION.value, + stage=GuardrailEventHooks.post_call if input_type == "response" else GuardrailEventHooks.pre_call, + ) # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) 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 b73d3adb99e..3bc0dfabefc 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 @@ -721,10 +721,18 @@ class PanwPrismaAirsHandler(CustomGuardrail): } } - def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None: + def _record_scan_id( + self, request_data: dict[str, object], scan_result: Mapping[str, object], stage: GuardrailEventHooks + ) -> None: """Surface the AIRS scan id on the response, so allowed calls are auditable too.""" scan_id: Final = scan_result.get("scan_id") - add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None) + add_guardrail_scan_id( + request_data=request_data, + scan_id=str(scan_id) if scan_id else None, + guardrail_name=self.guardrail_name, + provider=self._PROVIDER_NAME, + stage=stage, + ) def _handle_api_error_with_logging( self, @@ -948,7 +956,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): event_type=GuardrailEventHooks.post_call, ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool: """ @@ -1078,7 +1086,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.pre_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.pre_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1199,7 +1207,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.post_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1401,7 +1409,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) # Add guardrail to applied guardrails header for observability add_guardrail_to_applied_guardrails_header( @@ -1475,7 +1483,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_args = self._masked_tool_call_arguments( @@ -1829,7 +1841,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): new_texts.append(text) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_text = self._get_masked_text(scan_result, is_response=is_response) @@ -1901,7 +1917,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) # If we reach here, fallback_on_error="allow" else: - self._record_scan_id(request_data, mcp_scan_result) + self._record_scan_id(request_data, mcp_scan_result, GuardrailEventHooks.pre_call) action = mcp_scan_result.get("action", "block") masked_text = self._get_masked_text(mcp_scan_result, is_response=False) if action == "allow": diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 924f84be5f4..3250ae5cca9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -235,6 +235,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, "pillar_response_headers", @@ -291,6 +292,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 66f77db6da9..ecb2375d495 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,30 +1,33 @@ import copy +import json import sys from types import ModuleType, SimpleNamespace +from typing import Final +from unittest.mock import patch import pytest - +import litellm +from litellm.caching.caching import DualCache +from litellm.constants import MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( + _serialize_scan_metadata_header, add_guardrail_scan_id, add_policy_to_applied_policies_header, decrypt_callback_vars, encrypt_callback_vars, get_logging_caching_headers, - initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, + initialize_callbacks_on_proxy, normalize_callback_names, + process_callback, sanitize_openai_provider_metadata, strip_callback_config, ) -import litellm -from litellm.caching.caching import DualCache -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging - -from unittest.mock import patch -from litellm.proxy.common_utils.callback_utils import process_callback +from litellm.types.guardrails import GuardrailEventHooks def test_get_remaining_tokens_and_requests_from_request_data(): @@ -189,20 +192,109 @@ def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata(): assert headers["x-litellm-policy-sources"] == "global-baseline=team_default" +def _record( + request_data: dict[str, object], + scan_id: str | None, + guardrail_name: str = "airs", + provider: str = "panw_prisma_airs", + stage: GuardrailEventHooks = GuardrailEventHooks.pre_call, +) -> None: + add_guardrail_scan_id( + request_data=request_data, scan_id=scan_id, guardrail_name=guardrail_name, provider=provider, stage=stage + ) + + def test_add_guardrail_scan_id_dedupes_and_becomes_response_header(): request_data = {"litellm_metadata": {}} - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-2") - add_guardrail_scan_id(request_data=request_data, scan_id=None) + _record(request_data, "scan-1") + _record(request_data, "scan-1") + _record(request_data, "scan-2") + _record(request_data, None) assert request_data["litellm_metadata"]["guardrail_scan_ids"] == ("scan-1", "scan-2") assert get_logging_caching_headers(request_data)["x-litellm-guardrail-scan-id"] == "scan-1,scan-2" -def test_get_logging_caching_headers_omits_scan_id_header_without_scans(): - assert "x-litellm-guardrail-scan-id" not in get_logging_caching_headers({"litellm_metadata": {}}) +def test_scan_metadata_header_maps_each_id_to_its_guardrail_stage_and_provider(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + + _record( + request_data, "scan-1", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "mod-1", guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record(request_data, None, guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1,mod-1,scan-2" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "airs", "stage": "pre_call", "provider": "panw_prisma_airs", "scan_id": "scan-1"}, + {"guardrail": "mod", "stage": "pre_call", "provider": "openai_moderation", "scan_id": "mod-1"}, + {"guardrail": "airs", "stage": "post_call", "provider": "panw_prisma_airs", "scan_id": "scan-2"}, + ] + + +def test_scan_metadata_keeps_same_id_reused_across_stages(): + request_data: Final[dict[str, object]] = {"metadata": {}} + + _record(request_data, "scan-1", stage=GuardrailEventHooks.pre_call) + _record(request_data, "scan-1", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + assert [entry["stage"] for entry in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + "pre_call", + "post_call", + ] + + +def test_scan_metadata_header_drops_trailing_entries_to_stay_within_length_limit(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + scan_ids: Final = tuple(f"0f9c4b7e-3d2a-4c1b-9e8f-{index:012d}" for index in range(40)) + for scan_id in scan_ids: + _record(request_data, scan_id, stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == ",".join(scan_ids) + header: Final = headers["x-litellm-guardrail-scan-metadata"] + assert len(header) <= MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH + kept: Final = json.loads(header) + assert 1 < len(kept) < len(scan_ids) + assert [entry["scan_id"] for entry in kept] == list(scan_ids[: len(kept)]) + + +def test_serialize_scan_metadata_header_keeps_exactly_the_entries_that_fit(): + entries: Final = ({"scan_id": "a"}, {"scan_id": "b"}, {"scan_id": "c"}) + two_entries: Final = '[{"scan_id":"a"},{"scan_id":"b"}]' + + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries)) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) - 1) == '[{"scan_id":"a"}]' + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) + 1) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=1000) == json.dumps(entries, separators=(",", ":")) + assert _serialize_scan_metadata_header(entries, max_length=5) is None + assert _serialize_scan_metadata_header((), max_length=1000) is None + + +def test_scan_metadata_is_an_internal_metadata_key(): + assert sanitize_openai_provider_metadata({"guardrail_scan_metadata": "x", "keep": "y"}) == {"keep": "y"} + + +def test_get_logging_caching_headers_omits_scan_headers_without_scans(): + headers: Final = get_logging_caching_headers({"litellm_metadata": {}}) + assert headers is not None + assert "x-litellm-guardrail-scan-id" not in headers + assert "x-litellm-guardrail-scan-metadata" not in headers def test_initialize_callbacks_on_proxy_instantiates_compression_interception( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 2b43720a126..615d06b0f42 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -3,14 +3,19 @@ Test OpenAI Moderation Guardrail """ +import json import os +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) @@ -989,3 +994,30 @@ async def test_openai_moderation_initialize_guardrail_forwards_streaming_flags() assert guardrail.streaming_sampling_rate == 2 finally: litellm.logging_callback_manager._reset_all_callbacks() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("input_type", "stage"), [("request", "pre_call"), ("response", "post_call")]) +async def test_openai_moderation_records_moderation_id_as_scan_metadata(input_type: str, stage: str): + """Each moderation call's id is exposed with the guardrail name, stage and provider that produced it.""" + payload: Final = { + "id": f"modr-{stage}", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}, "category_applied_input_types": {}}], + } + http_client: Final = AsyncHTTPHandler() + http_client.client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail: Final = OpenAIModerationGuardrail(guardrail_name="openai-mod") + guardrail.async_handler = http_client + request_data: Final[dict[str, object]] = {"metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data=request_data, input_type=input_type) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == f"modr-{stage}" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "openai-mod", "stage": stage, "provider": "openai_moderation", "scan_id": f"modr-{stage}"} + ] 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 8f29ba66814..3d7c6e06d94 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 @@ -12,6 +12,7 @@ This test file follows LiteLLM's testing patterns and covers: import copy import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5785,7 +5786,14 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123" - assert "x-litellm-guardrail-scan-metadata" not in headers + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + { + "guardrail": handler.guardrail_name, + "stage": "pre_call", + "provider": "panw_prisma_airs", + "scan_id": "scan-abc-123", + } + ] @pytest.mark.asyncio async def test_request_and_response_scan_ids_are_both_exposed(self, user_api_key_dict): @@ -5809,6 +5817,26 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123,scan-response-456" + assert [(e["stage"], e["scan_id"]) for e in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + ("pre_call", "scan-abc-123"), + ("post_call", "scan-response-456"), + ] + + @pytest.mark.asyncio + async def test_apply_guardrail_response_scan_is_tagged_post_call(self): + from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers + + handler: Final = self._handler(self.ALLOW_SCAN_RESULT) + request_data: Final[dict[str, object]] = {"litellm_call_id": "test-call-id", "model": "gpt-4", "metadata": {}} + + await handler.apply_guardrail( + inputs={"texts": ["Hello world"]}, request_data=request_data, input_type="response" + ) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + entries: Final = json.loads(headers["x-litellm-guardrail-scan-metadata"]) + assert [(e["stage"], e["provider"]) for e in entries] == [("post_call", "panw_prisma_airs")] @pytest.mark.asyncio async def test_repeated_scan_id_is_not_duplicated(self, user_api_key_dict): @@ -5850,6 +5878,8 @@ class TestPanwAirsScanIdExposure: assert "guardrail_scan_ids" in _UNTRUSTED_METADATA_CONTROL_FIELDS assert "guardrail_scan_ids" in _UNTRUSTED_ROOT_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_METADATA_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_ROOT_CONTROL_FIELDS class TestPanwAirsBlockedErrorDetailPassthrough: """Regression tests for the full AIRS scan response on blocks. From 0c0a1dd76f31b56059508dd674fee1764c83828f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 06:34:36 +0000 Subject: [PATCH 091/157] test(cli): avoid mutable request tracking Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/test_global_options.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 6de56573ffb..9ed9f91ab9f 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -37,16 +37,17 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output -def test_lite_version_does_not_fetch_model_cost_map(): +def test_lite_version_does_not_fetch_model_cost_map(tmp_path: Path) -> None: lite_path = shutil.which("lite") if lite_path is None: pytest.skip("lite executable is unavailable") - requests: list[str] = [] + request_log: Path = tmp_path / "requests.log" class _CostMapHandler(BaseHTTPRequestHandler): def do_GET(self) -> None: - requests.append(self.path) + with request_log.open("a", encoding="utf-8") as log_file: + log_file.write(f"{self.path}\n") body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' self.send_response(200) self.send_header("Content-Type", "application/json") @@ -84,7 +85,8 @@ def test_lite_version_does_not_fetch_model_cost_map(): assert result.returncode == 0 assert "LiteLLM Proxy CLI Version" in result.stdout - assert requests == [] + request_count: int = request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 + assert request_count == 0 def test_cli_source_is_ascii_only(): From 9d90a544916c6f38397862dd18b6a4c12de98827 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 9 Sep 2026 00:06:26 -0700 Subject: [PATCH 092/157] revert(model-management): roll back #40047 This reverts commit e8e3172d7d70558929f32f057ebe4c7471c8c352 Restore the previous model update and router cost registration behavior while pricing compatibility is investigated --- .../model_management_endpoints.py | 70 +--- litellm/router.py | 26 +- tests/e2e/coverage_registry/mgmt.yaml | 2 - .../management/test_model_lifecycle_e2e.py | 365 ------------------ tests/e2e/models.py | 74 +--- tests/e2e/proxy_client.py | 199 +--------- tests/e2e/test_proxy_client.py | 57 +-- tests/e2e/transport.py | 1 - .../test_model_management_endpoints.py | 178 +-------- .../test_router_model_cost_isolation.py | 192 --------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 11 files changed, 39 insertions(+), 1130 deletions(-) delete mode 100644 tests/e2e/management/test_model_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 742b9d9817f..0f19e9ce149 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -119,7 +119,6 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, - LiteLLM_Params, ModelInfo, updateDeployment, ) @@ -729,44 +728,6 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) -_OWNERSHIP_FIELDS: Final = frozenset( - { - "db_model", - "team_id", - "team_public_model_name", - "access_groups", - "created_at", - "created_by", - "updated_at", - "updated_by", - "blocked", - } -) - -_STORED_REQUIRED_FIELDS: Final = frozenset( - name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() -) - -_NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS | frozenset(PTU_MODEL_INFO_FIELDS) - - -def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: - """The keys a patch sends as an explicit null, which update_db_model removes from the - stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored - models require, since clearing one writes a row no reload can rebuild through - LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, - whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob - the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. - """ - if patch is None: - return frozenset() - return frozenset( - field - for field in patch.model_fields_set - if field not in _NULL_CLEAR_IGNORED_FIELDS and getattr(patch, field) is None - ) - - def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -787,15 +748,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.model_info: merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) - for field in _explicitly_cleared_fields(updated_patch.litellm_params): - merged_litellm_params.pop(field, None) - if field in SPECIAL_MODEL_INFO_PARAMS: - merged_model_info.pop(field, None) - for field in _explicitly_cleared_fields(updated_patch.model_info): - merged_model_info.pop(field, None) - if field in SPECIAL_MODEL_INFO_PARAMS: - merged_litellm_params.pop(field, None) + # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI + # passes through (which today re-sends the OLD pricing on every save) cannot + # silently undo a litellm_params clear via .update(). + # + # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character + # and cache read/write costs) so this path cannot be used to null out privileged + # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are + # mirrored between litellm_params and model_info by Deployment.__init__, so the + # clear propagates to both blobs. + if updated_patch.litellm_params: + for field in updated_patch.litellm_params.model_fields_set: + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: + for field in updated_patch.model_info.model_fields_set: + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) @@ -845,9 +816,8 @@ async def patch_model( """ PATCH Endpoint for partial model updates. - JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the - body is unchanged, a key sent as null is removed from the stored row, and a key sent with a - value is set (identity and ownership keys such as `id` and `team_id` ignore a null). + Only updates the fields specified in the request while preserving other existing values. + Follows proper PATCH semantics by only modifying provided fields. Args: model_id: The ID of the model to update diff --git a/litellm/router.py b/litellm/router.py index 1252d7e7487..934a4ac86a9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -628,15 +628,6 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 -# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat -# namespace with the built-in model catalog. Only a key it created may be evicted, or a -# deployment whose id names a real model would strip that model's pricing and -# capabilities for every other deployment of it. delete_deployment gives a key back once no -# live router still serves that id, so a later catalog refresh that starts serving the name -# is not treated as a deployment's own. -_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys - - class FallbackAwareStreamWrapper(CustomStreamWrapper): """Base for the Router's chat-completion stream wrappers, which are built around the attempt the Router picked first and have to repoint themselves when a fallback takes over.""" @@ -9770,7 +9761,6 @@ class Router: """ idx: Final = len(self.model_list) self.model_list.append(model) - _live_routers.add(self) # mutable-ok: track dynamic routers without extending their lifetimes self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() @@ -9939,12 +9929,7 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. An - entry this function created is replaced rather than merged, so a price cleared - from the deployment does not linger from an earlier registration and keep - billing at the old rate. An entry it did not create is left to merge, because - a deployment id that collides with a catalog model name shares that model's - entry with every other deployment of it. + the entries a refresh rebuilds are the ones a fresh boot would produce. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9961,10 +9946,6 @@ class Router: } if model_id is not None: - if model_id in _DEPLOYMENT_COST_MAP_KEYS: - litellm.model_cost.pop(model_id, None) # mutable-ok: remove cleared prices from the shared entry - elif model_id not in litellm.model_cost: - _DEPLOYMENT_COST_MAP_KEYS.add(model_id) # mutable-ok: retain shared ownership across reloads litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, @@ -10061,11 +10042,6 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) - if not any( - router is not self and id in router.model_id_to_deployment_index_map - for router in tuple(_live_routers) - ): - _DEPLOYMENT_COST_MAP_KEYS.discard(id) # mutable-ok: the last owning router released this key try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 83f1711a245..c8d7037d2fd 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,8 +76,6 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} -- {id: mgmt.model.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "model_management_endpoints.py:731", rationale: "A partial PATCH changes only the keys it names; every other stored key reads back byte-for-byte, and the new rate reaches billing"} -- {id: mgmt.model.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "model_management_endpoints.py:731", fail_before_fix: proven, rationale: "An explicit null on PATCH removes the key from the stored row and billing falls back to the cost map; before the fix only mirrored pricing keys could be cleared"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/management/test_model_lifecycle_e2e.py b/tests/e2e/management/test_model_lifecycle_e2e.py deleted file mode 100644 index 99044efeb46..00000000000 --- a/tests/e2e/management/test_model_lifecycle_e2e.py +++ /dev/null @@ -1,365 +0,0 @@ -"""Live e2e: the lifecycle of a DB-stored deployment through the model management -routes, read back on every gateway replica. - -Each test registers its own gpt-4o-mini mock deployment through /model/new (deleted -on teardown) with non-default pricing, context window, mode, and api_base pinned, then -walks the lifecycle up to the step it proves: the create reads back field for field, -a partial PATCH changes only the key it names, an explicit null on PATCH removes the -key from the stored row (JSON Merge Patch), a call after the price clear is billed at -the cost map's rate rather than the cleared override, and a delete removes the -deployment from /model/info and makes the model name unknown to /chat/completions. - -The stored row is read back from /model/info, a control-plane route with one answer -behind it. What every gateway must agree on is which models it serves, so the create -and delete steps poll /v1/models on every URL in PROXY_REPLICA_URLS through -ProxyClient.read_model_back_everywhere, failing by name on the gateway that never converged. -""" - -from __future__ import annotations - -import math -import time -from collections.abc import Callable -from dataclasses import dataclass -from typing import Final - -import pytest -from pydantic import BaseModel - -from e2e_config import unique_marker -from e2e_http import unwrap -from lifecycle import ResourceManager -from management_client import ManagementClient -from models import ( - ChatBody, - ChatMessage, - Clear, - LiteLLMParamsBody, - LiteLLMParamsPatch, - ModelInfoBody, - ModelInfoEntry, - ModelInfoResponse, - ModelNewBody, - ModelPatchBody, - ModelsListResponse, - SpendLogRow, -) - -pytestmark = pytest.mark.e2e - -BACKEND_MODEL: Final = "gpt-4o-mini" -PINNED_API_BASE: Final = "https://pinned.example.invalid/v1" -PINNED_MAX_INPUT_TOKENS: Final = 4096 -PINNED_INPUT_RATE: Final = 1e-05 -UPDATED_INPUT_RATE: Final = 2e-05 -PINNED_OUTPUT_RATE: Final = 3e-05 - -# A PATCH lands on the control plane, and each gateway picks it up on its own config -# reload, so the first call after the write can still be billed at the old rate. There -# is no price on the gateway's data-plane surface to poll, so the billing steps drive -# calls until the new rate shows up in the spend row and let the deadline be what fails. -BILLING_CONVERGENCE_TIMEOUT: Final = 90.0 -BILLING_CONVERGENCE_INTERVAL: Final = 5.0 - - -@dataclass(frozen=True, slots=True) -class Registered: - model_name: str - model_id: str - - -class _ErrorDetail(BaseModel): - message: str - - -class _ErrorEnvelope(BaseModel): - error: _ErrorDetail - - -def _register(client: ManagementClient, resources: ResourceManager) -> Registered: - """Register a mock gpt-4o-mini deployment with every field under test pinned to a - non-default value, deleted on teardown. max_input_tokens is pinned in - litellm_params only: a value in model_info is copied into the shared cost-map - entry for the backend model, which would leak into every other gpt-4o-mini - deployment on the proxy.""" - model_name: Final = f"e2e-lifecycle-{unique_marker()}" - model_id: Final = client.proxy.register_model( - ModelNewBody( - model_name=model_name, - litellm_params=LiteLLMParamsBody( - model=BACKEND_MODEL, - mock_response="ok", - api_base=PINNED_API_BASE, - input_cost_per_token=PINNED_INPUT_RATE, - output_cost_per_token=PINNED_OUTPUT_RATE, - max_input_tokens=PINNED_MAX_INPUT_TOKENS, - ), - model_info=ModelInfoBody(mode="chat"), - ) - ) - resources.defer(lambda: client.proxy.delete_model(model_id)) - return Registered(model_name=model_name, model_id=model_id) - - -def _entry(body: ModelInfoResponse, model_name: str) -> ModelInfoEntry | None: - return next((entry for entry in body.data if entry.model_name == model_name), None) - - -def _stored_entry( - client: ManagementClient, - model_name: str, - *, - converged: Callable[[ModelInfoEntry], bool], -) -> ModelInfoEntry: - """The stored /model/info row for `model_name`, once it satisfies `converged`. - - /model/info is a control-plane route: the gateways named in PROXY_REPLICA_URLS - serve the LLM surface only, so the stored row has one answer, not one per - gateway. What every gateway must agree on is which models it serves, and - `_assert_served_everywhere` / `_assert_absent_everywhere` poll /v1/models for - that.""" - - def has_converged(body: ModelInfoResponse) -> bool: - entry: Final = _entry(body, model_name) - return entry is not None and converged(entry) - - body: Final = client.proxy.read_model_back("/model/info", ModelInfoResponse, predicate=has_converged) - entry: Final = _entry(body, model_name) - assert entry is not None, f"/model/info stopped listing {model_name!r} between the poll and the read" - return entry - - -def _serves(body: ModelsListResponse, model_name: str) -> bool: - return any(entry.id == model_name for entry in body.data) - - -def _assert_served_everywhere(client: ManagementClient, model_name: str) -> None: - _ = client.proxy.read_model_back_everywhere( - "/v1/models", ModelsListResponse, predicate=lambda body: _serves(body, model_name) - ) - - -def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: - _ = client.proxy.read_model_back_everywhere( - "/v1/models", ModelsListResponse, predicate=lambda body: not _serves(body, model_name) - ) - - -def _assert_untouched_keys_as_created(entry: ModelInfoEntry, replica: str) -> None: - """The keys no later step names read back byte-for-byte as /model/new wrote them.""" - params: Final = entry.litellm_params - assert params.model == BACKEND_MODEL, f"{replica}: litellm_params.model {params.model!r} != {BACKEND_MODEL!r}" - assert params.api_base == PINNED_API_BASE, f"{replica}: api_base {params.api_base!r} != {PINNED_API_BASE!r}" - assert params.output_cost_per_token == PINNED_OUTPUT_RATE, ( - f"{replica}: output_cost_per_token {params.output_cost_per_token} != {PINNED_OUTPUT_RATE}" - ) - assert entry.model_info.mode == "chat", f"{replica}: model_info.mode {entry.model_info.mode!r} != 'chat'" - - -def _approx_equal(actual: float, expected: float) -> bool: - return math.isclose(actual, expected, rel_tol=1e-2, abs_tol=1e-9) - - -def _priced(rows: list[SpendLogRow]) -> bool: - return any(row.metadata and row.metadata.cost_breakdown and row.metadata.cost_breakdown.input_cost for row in rows) - - -def _billed_input_cost(client: ManagementClient, model_name: str, key: str) -> tuple[int, float]: - """Drive one chat completion through `model_name` and return the prompt tokens and - input cost its spend row recorded, so a test can assert the rate the gateway actually - billed rather than only the rate it stored.""" - chat: Final = unwrap( - client.proxy.chat( - key, - ChatBody( - model=model_name, - messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], - max_tokens=16, - ), - ) - ) - assert chat.id is not None, f"chat completion carried no id to find its spend row by: {chat}" - - rows: Final = client.proxy.poll_logs_for_request_id(chat.id, predicate=_priced) - row: Final = next((row for row in rows if row.request_id == chat.id), None) - assert row is not None and row.metadata and row.metadata.cost_breakdown, ( - f"no priced spend row for request {chat.id} before the deadline: {rows}" - ) - prompt_tokens: Final = row.prompt_tokens or 0 - input_cost: Final = row.metadata.cost_breakdown.input_cost - assert prompt_tokens > 0 and input_cost is not None, f"spend row logged no prompt tokens or input cost: {row}" - return prompt_tokens, input_cost - - -def _await_billed_input_cost( - client: ManagementClient, model_name: str, key: str, *, expected_rate: float -) -> tuple[int, float]: - """Drive calls through `model_name` until one is billed at `expected_rate`, and - return the prompt tokens and input cost of the last spend row either way. - - Only the deadline ends the wait unsatisfied: a rate that never reaches the gateway - comes back as the stale cost for the caller to assert on, so the rate the caller - expects is still what decides the test.""" - deadline: Final = time.monotonic() + BILLING_CONVERGENCE_TIMEOUT - while True: - prompt_tokens, input_cost = _billed_input_cost(client, model_name, key) - if _approx_equal(input_cost, prompt_tokens * expected_rate) or time.monotonic() >= deadline: - return prompt_tokens, input_cost - time.sleep(BILLING_CONVERGENCE_INTERVAL) - - -class TestModelLifecycle: - @pytest.mark.covers("mgmt.model.add.persists") - def test_create_reads_back_every_field_and_serves_on_every_replica( - self, client: ManagementClient, resources: ResourceManager - ) -> None: - registered = _register(client, resources) - - entry = _stored_entry(client, registered.model_name, converged=lambda _entry: True) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( - f"{stored}: input_cost_per_token {entry.litellm_params.input_cost_per_token} != {PINNED_INPUT_RATE}" - ) - assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( - f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" - ) - assert entry.model_info.id == registered.model_id, ( - f"{stored}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" - ) - - _assert_served_everywhere(client, registered.model_name) - - @pytest.mark.covers("mgmt.model.update.preserves_unrelated_fields") - def test_partial_update_changes_only_the_named_key( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - - stored = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(input_cost_per_token=UPDATED_INPUT_RATE)), - ) - assert stored.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, ( - f"PATCH response stores input_cost_per_token {stored.litellm_params.input_cost_per_token}, " - f"sent {UPDATED_INPUT_RATE}" - ) - - entry = _stored_entry( - client, - registered.model_name, - converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, - ) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( - f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" - ) - assert entry.model_info.input_cost_per_token == UPDATED_INPUT_RATE, ( - f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} " - f"did not mirror the updated {UPDATED_INPUT_RATE}" - ) - - prompt_tokens, input_cost = _await_billed_input_cost( - client, registered.model_name, scoped_key, expected_rate=UPDATED_INPUT_RATE - ) - assert _approx_equal(input_cost, prompt_tokens * UPDATED_INPUT_RATE), ( - f"input_cost {input_cost} != {prompt_tokens} tokens * updated rate {UPDATED_INPUT_RATE} " - f"= {prompt_tokens * UPDATED_INPUT_RATE}; the partial update did not reach billing" - ) - - @pytest.mark.covers("mgmt.model.update.clear_persists") - def test_explicit_null_removes_the_key_from_the_stored_row( - self, client: ManagementClient, resources: ResourceManager - ) -> None: - registered = _register(client, resources) - - stored = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), - ) - stored_params = stored.litellm_params.model_fields_set - assert "max_input_tokens" not in stored_params, ( - f"stored litellm_params still carries max_input_tokens " - f"{stored.litellm_params.max_input_tokens} after an explicit null" - ) - assert "input_cost_per_token" not in stored_params, ( - f"stored litellm_params still carries input_cost_per_token " - f"{stored.litellm_params.input_cost_per_token} after an explicit null" - ) - assert "max_input_tokens" not in stored.model_info.model_fields_set, ( - f"stored model_info carries max_input_tokens {stored.model_info.max_input_tokens} after the clear" - ) - assert "input_cost_per_token" not in stored.model_info.model_fields_set, ( - f"stored model_info still mirrors input_cost_per_token {stored.model_info.input_cost_per_token}" - ) - - cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token - assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" - entry = _stored_entry( - client, - registered.model_name, - converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, - ) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - served = entry.litellm_params.model_fields_set - assert "max_input_tokens" not in served, ( - f"{stored}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" - ) - assert "input_cost_per_token" not in served, ( - f"{stored}: litellm_params still serves input_cost_per_token {entry.litellm_params.input_cost_per_token}" - ) - assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( - f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} is not the " - f"cost map's {cost_map_input_rate}; the cleared override {PINNED_INPUT_RATE} still resolves" - ) - - @pytest.mark.covers("mgmt.model.update.clear_persists") - def test_cleared_price_is_billed_at_the_cost_map_rate( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - _ = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), - ) - _ = _stored_entry( - client, - registered.model_name, - converged=lambda entry: "input_cost_per_token" not in entry.litellm_params.model_fields_set, - ) - cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token - assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" - - prompt_tokens, input_cost = _await_billed_input_cost( - client, registered.model_name, scoped_key, expected_rate=cost_map_input_rate - ) - - assert _approx_equal(input_cost, prompt_tokens * cost_map_input_rate), ( - f"input_cost {input_cost} != {prompt_tokens} tokens * cost map rate {cost_map_input_rate} " - f"= {prompt_tokens * cost_map_input_rate}" - ) - assert not _approx_equal(input_cost, prompt_tokens * PINNED_INPUT_RATE), ( - f"input_cost {input_cost} is still billed at the cleared override {PINNED_INPUT_RATE}" - ) - - @pytest.mark.covers("mgmt.model.delete.persists") - def test_delete_removes_the_deployment_everywhere( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - _ = _stored_entry(client, registered.model_name, converged=lambda _entry: True) - - client.delete_model_strict(registered.model_id) - - _assert_absent_everywhere(client, registered.model_name) - refused = client.chat_status(scoped_key, registered.model_name, "hi this is a test") - assert refused.status_code == 400, ( - f"chat against the deleted model must be rejected 400, got {refused.status_code}: {refused.body[:300]}" - ) - envelope = _ErrorEnvelope.model_validate_json(refused.body) - assert envelope.error.message, f"400 body must carry an error message: {refused.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8db37bd25a5..62810e6cfd9 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,11 +655,6 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- -class CostBreakdown(BaseModel): - input_cost: float | None = None - output_cost: float | None = None - - class GuardrailEntityMatch(BaseModel): entity_type: str score: float @@ -677,7 +672,6 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): - cost_breakdown: CostBreakdown | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None @@ -821,42 +815,15 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token -class DeploymentParams(CustomPricing): - """The litellm_params half of a /model/info row: the stored deployment as written, - credentials scrubbed. Unlike model_info it is never back-filled from the cost map, - so a key the store dropped is absent here (check `model_fields_set`).""" - - model: str | None = None - api_base: str | None = None - max_input_tokens: int | None = None - - -class DeploymentModelInfo(CustomPricing): - id: str | None = None - max_input_tokens: int | None = None - - class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for - it - the override merged over the cost-map defaults, so a key cleared from the - stored blob reads as the cost-map default here.""" + it - the override merged over the cost-map defaults.""" model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: DeploymentParams = DeploymentParams() - model_info: DeploymentModelInfo = DeploymentModelInfo() - - -class StoredDeployment(BaseModel): - """PATCH /model/{model_id}/update answers with the row as stored: both blobs raw, - nothing back-filled, so a cleared key is absent from `model_fields_set` of the - blob it was cleared from.""" - - model_config = ConfigDict(protected_namespaces=()) - model_name: str - litellm_params: DeploymentParams - model_info: DeploymentModelInfo + litellm_params: CustomPricing = CustomPricing() + model_info: CustomPricing = CustomPricing() class ModelInfoResponse(BaseModel): @@ -957,10 +924,9 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None - max_input_tokens: int | None = None -ModelMode = Literal["chat", "batch", "realtime", "image_generation"] +ModelMode = Literal["batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): @@ -970,7 +936,6 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None - max_input_tokens: int | None = None access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None @@ -999,37 +964,6 @@ class ModelUpdateBody(BaseModel): model_info: ModelInfoBody -class Clear(BaseModel): - """Serializes to JSON null. The transport dumps every body with exclude_none, so a - field set to this is how a patch carries the explicit null that removes a stored key.""" - - @model_serializer - def _as_null(self) -> None: - return None - - -class LiteLLMParamsPatch(BaseModel): - api_base: str | Clear | None = None - max_input_tokens: int | Clear | None = None - input_cost_per_token: float | Clear | None = None - output_cost_per_token: float | Clear | None = None - - -class ModelInfoPatch(BaseModel): - mode: ModelMode | Clear | None = None - max_input_tokens: int | Clear | None = None - - -class ModelPatchBody(BaseModel): - """PATCH /model/{model_id}/update body, JSON Merge Patch over the stored deployment: - a field left None is dropped from the body and unchanged, a field set to `Clear()` - is sent as null and removed, a field with a value is set.""" - - model_config = ConfigDict(protected_namespaces=()) - litellm_params: LiteLLMParamsPatch | None = None - model_info: ModelInfoPatch | None = None - - class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index fa1b06fe7ed..1bac5116a9d 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,10 +10,10 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass -from datetime import datetime from functools import reduce +from datetime import datetime from types import MappingProxyType from typing import Final @@ -62,7 +62,6 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, - ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -73,7 +72,6 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, - StoredDeployment, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -134,103 +132,6 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None -type BodyReader[R: BaseModel] = Callable[[float], Result[R]] - - -@dataclass(frozen=True, slots=True) -class BodyNotConverged[R: BaseModel]: - """The deadline passed without a read the predicate accepted; `last_result` is the - final read, so the caller can tell a body that never matched from a read that - failed.""" - - last_result: Result[R] | None - - -@dataclass(frozen=True, slots=True) -class BodyConverged[R: BaseModel]: - """Every replica answered a body the predicate accepted; `bodies` is the last read - per replica.""" - - bodies: Mapping[str, R] - - -@dataclass(frozen=True, slots=True) -class BodyNeverConvergedOn[R: BaseModel]: - """`BodyNotConverged` labeled with the replica whose reads never satisfied the predicate.""" - - replica: str - last_result: Result[R] | None - - -def await_body_converged[R: BaseModel]( - read: BodyReader[R], - *, - predicate: Callable[[R], bool], - timeout: float, - interval: float, - request_timeout: float, - now: Callable[[], float], - sleep: Callable[[float], None], -) -> Success[R] | BodyNotConverged[R]: - """Poll `read` until it answers a body `predicate` accepts, or `timeout` passes. - - Each read's request timeout is clamped to the remaining budget, and the sleep - between reads to the time left, so the last read before the deadline is never - skipped. Clock and sleep are injected.""" - deadline: Final = now() + timeout - - def reads() -> Iterator[Result[R]]: - while (remaining := deadline - now()) > 0: - yield read(min(request_timeout, remaining)) - sleep(min(interval, max(deadline - now(), 0.0))) - - def attempts() -> Iterator[Success[R] | BodyNotConverged[R]]: - for result in reads(): - if isinstance(result, Success) and predicate(result.data): - yield result - return - yield BodyNotConverged(last_result=result) - - initial: Final[Success[R] | BodyNotConverged[R]] = BodyNotConverged(last_result=None) - return reduce(lambda _previous, result: result, attempts(), initial) - - -def await_body_converged_everywhere[R: BaseModel]( - readers: Mapping[str, BodyReader[R]], - *, - predicate: Callable[[R], bool], - timeout: float, - interval: float, - request_timeout: float, - now: Callable[[], float], - sleep: Callable[[float], None], -) -> BodyConverged[R] | BodyNeverConvergedOn[R]: - """`await_body_converged` against every replica in turn, each with the full budget, so a - write counts as landed only once every replica serves it.""" - def read_replica( - outcome: BodyConverged[R] | BodyNeverConvergedOn[R], - item: tuple[str, BodyReader[R]], - ) -> BodyConverged[R] | BodyNeverConvergedOn[R]: - if isinstance(outcome, BodyNeverConvergedOn): - return outcome - replica, read = item - match await_body_converged( - read, - predicate=predicate, - timeout=timeout, - interval=interval, - request_timeout=request_timeout, - now=now, - sleep=sleep, - ): - case Success(data=data): - return BodyConverged(bodies=MappingProxyType({**outcome.bodies, replica: data})) - case BodyNotConverged(last_result=last_result): - return BodyNeverConvergedOn(replica=replica, last_result=last_result) - initial: Final[BodyConverged[R] | BodyNeverConvergedOn[R]] = BodyConverged(bodies=MappingProxyType({})) - return reduce(read_replica, readers.items(), initial) - - def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -757,102 +658,6 @@ class ProxyClient: ) ) - def patch_model(self, model_id: str, body: ModelPatchBody) -> StoredDeployment: - """JSON Merge Patch the deployment `model_id` via PATCH /model/{model_id}/update: - a field the body omits is unchanged, one sent as null is removed from the stored - row, one sent with a value is set. See ModelPatchBody for how a null is sent. - Returns the row as stored after the write.""" - return unwrap( - self.transport.patch( - f"/model/{model_id}/update", - headers=self.transport.master, - json=body, - response_type=StoredDeployment, - ) - ) - - def read_model_back_everywhere[R: BaseModel]( - self, path: str, response_type: type[R], *, predicate: Callable[[R], bool] - ) -> Mapping[str, R]: - """GET `path` on every replica until each answers a body `predicate` accepts, - polling to poll_timeout, and return the last body per replica. - - Fails naming the replica that never converged, so a write that reached one - gateway but not the others is caught instead of passing on whichever gateway - the balancer answered from. Falls back to the single proxy address when no - replica list is configured. - - `path` must be a data-plane route. The replicas are gateways, which serve only - the LLM surface, so a control-plane path answers on exactly one service and - 404s on every replica in a split deployment: asking each replica for one is - never the question the caller means. Read those through `self.transport` - instead, which routes them to the control plane.""" - if is_control_plane_path(path): - raise AssertionError( - f"read_model_back_everywhere({path!r}) asks every data-plane replica for a control-plane route. " - "The replicas are gateways and do not serve it; poll a data-plane path such as /v1/models " - "here, and read the control plane through the shared transport." - ) - readers: Final = { - url: self._body_reader(transport, path, response_type) - for url, transport in self._read_back_replicas().items() - } - outcome: Final = await_body_converged_everywhere( - readers, - predicate=predicate, - timeout=self.poll_timeout, - interval=self.poll_interval, - request_timeout=REQUEST_TIMEOUT, - now=time.monotonic, - sleep=time.sleep, - ) - match outcome: - case BodyConverged(bodies=bodies): - return bodies - case BodyNeverConvergedOn(replica=replica, last_result=last_result): - raise AssertionError( - f"GET {path} on {replica} never answered the expected body within " - f"{self.poll_timeout}s; last read: {last_result}" - ) - - def read_model_back[R: BaseModel](self, path: str, response_type: type[R], *, predicate: Callable[[R], bool]) -> R: - """GET `path` through the shared transport until the body satisfies `predicate`, - polling to poll_timeout, and return that body. - - The counterpart to `read_model_back_everywhere` for a control-plane route such as - /model/info: the stored row lives in one database behind one control plane, so - there is a single answer to converge on rather than one per gateway.""" - outcome: Final = await_body_converged_everywhere( - {CONTROL_PLANE_BASE_URL: self._body_reader(self.transport, path, response_type)}, - predicate=predicate, - timeout=self.poll_timeout, - interval=self.poll_interval, - request_timeout=REQUEST_TIMEOUT, - now=time.monotonic, - sleep=time.sleep, - ) - match outcome: - case BodyConverged(bodies=bodies): - return bodies[CONTROL_PLANE_BASE_URL] - case BodyNeverConvergedOn(last_result=last_result): - raise AssertionError( - f"GET {path} never answered the expected body within " - f"{self.poll_timeout}s; last read: {last_result}" - ) - - def _read_back_replicas(self) -> Mapping[str, Transport]: - return self.replicas or MappingProxyType({CONTROL_PLANE_BASE_URL: self.transport}) - - @staticmethod - def _body_reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> BodyReader[R]: - return lambda timeout: transport.get( - path, - headers=transport.master, - params=NoBody(), - response_type=response_type, - timeout=timeout, - ) - def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index cbf7f5648d4..3b84a47e3cc 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -20,12 +20,8 @@ from typing import Final, cast import pytest from e2e_config import parse_replica_urls from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( - BodyReader, - BodyConverged, - BodyNeverConvergedOn, - await_body_converged_everywhere, ConvergeOutcome, Converged, EverywhereConverged, @@ -278,54 +274,3 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") - - -def _info(*model_names: str) -> Success[ModelInfoResponse]: - entries: Final = [ModelInfoEntry(model_name=model_name) for model_name in model_names] - return Success(status_code=200, data=ModelInfoResponse(data=entries)) - - -def _reader(results: Iterable[Success[ModelInfoResponse]]) -> BodyReader[ModelInfoResponse]: - it: Final = iter(results) - return lambda _timeout: next(it) - - -def _lists_model(body: ModelInfoResponse) -> bool: - return any(entry.model_name == MODEL for entry in body.data) - - -def _read_back( - readers: Mapping[str, BodyReader[ModelInfoResponse]], -) -> tuple[BodyConverged[ModelInfoResponse] | BodyNeverConvergedOn[ModelInfoResponse], FakeClock]: - clock: Final = FakeClock() - outcome: Final = await_body_converged_everywhere( - readers, - predicate=_lists_model, - timeout=TIMEOUT, - interval=INTERVAL, - request_timeout=5.0, - now=clock.now, - sleep=clock.sleep, - ) - return outcome, clock - - -class TestAwaitBodyConvergedEverywhere: - def test_waits_for_the_lagging_replica_and_returns_every_body(self) -> None: - readers: Final = { - "gateway-1": _reader(repeat(_info(MODEL))), - "gateway-2": _reader(chain(repeat(_info(), 2), repeat(_info(MODEL)))), - } - outcome, clock = _read_back(readers) - assert outcome == BodyConverged(bodies={"gateway-1": _info(MODEL).data, "gateway-2": _info(MODEL).data}) - assert clock.elapsed == 2 * INTERVAL - - @pytest.mark.parametrize("lagging", ["gateway-1", "gateway-2"]) - def test_fails_naming_the_replica_that_never_converges(self, lagging: str) -> None: - readers: Final = { - "gateway-1": _reader(repeat(_info(MODEL))), - "gateway-2": _reader(repeat(_info(MODEL))), - } | {lagging: _reader(repeat(_info()))} - outcome, clock = _read_back(readers) - assert outcome == BodyNeverConvergedOn(replica=lagging, last_result=_info()) - assert clock.elapsed >= TIMEOUT diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 804e073a4a0..44fdbaa3e41 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,7 +306,6 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/config", "/guardrails", "/openapi.json", - "/public/", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 1376727e296..c02f886fc31 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3115,6 +3115,9 @@ class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored between the two by Deployment.__init__). + + Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) + cannot be cleared via this path. """ def test_clear_input_cost_removes_from_both_blobs(self): @@ -3190,10 +3193,10 @@ class TestUpdateDBModelClearPricing: assert params["input_cost_per_token"] == 0.000001 assert params["output_cost_per_token"] == 0.000007 - def test_null_on_one_field_leaves_other_fields_alone(self): - """A null clears only the key it names: pricing the patch never mentions and - the ownership key team_id stay put, so a team admin can't ungate a - team-scoped model through the clear path. + def test_null_on_non_pricing_field_does_not_clear(self): + """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. + Privileged or unrelated model_info fields (e.g. team_id) must be unaffected + by the null-clearing path so a team admin can't ungate a team-scoped model. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, @@ -3214,6 +3217,8 @@ class TestUpdateDBModelClearPricing: model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), ) + # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id + # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, updated_patch=updateDeployment( @@ -3389,171 +3394,6 @@ class TestUpdateDBModelClearPricing: assert info["cache_creation_input_token_cost"] == 0.000003 -_PROTECTED_MODEL_INFO_VALUES = { - "team_id": "team-keep-me", - "team_public_model_name": "team-facing-name", - "access_groups": ["group-a"], - "created_at": "2026-01-01T00:00:00+00:00", - "created_by": "creator", - "updated_at": "2026-01-02T00:00:00+00:00", - "updated_by": "updater", - "blocked": True, -} - - -def _build_db_model_with_pinned_model_info(): - """Deployment whose stored blobs pin non-pricing keys an earlier save wrote, next to a - pricing override, so a clear can be checked key by key.""" - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - - return Deployment( - model_name="pinned-gpt-4o-mini", - litellm_params=LiteLLM_Params( - model="gpt-4o-mini", input_cost_per_token=0.000001, max_input_tokens=4096 - ), - model_info=ModelInfo( - id="dep-pinned-0", - max_input_tokens=4096, - mode="chat", - supports_vision=True, - **_PROTECTED_MODEL_INFO_VALUES, - ), - ) - - -class TestUpdateDBModelNullClearsAnyKey: - """JSON Merge Patch on PATCH /model/{id}/update: a key sent as null is removed from the - stored blob it was sent in, whatever the key, except the identity and ownership keys, - whose nulls are ignored.""" - - def test_model_info_nulls_remove_pinned_non_pricing_keys(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"max_input_tokens": None, "mode": None}} - ), - ) - - info = json.loads(result["model_info"]) - assert "max_input_tokens" not in info - assert "mode" not in info - assert info["supports_vision"] is True - assert info["input_cost_per_token"] == 0.000001 - - def test_litellm_params_null_removes_pinned_non_pricing_key(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"litellm_params": {"max_input_tokens": None}} - ), - ) - - params = json.loads(result["litellm_params"]) - info = json.loads(result["model_info"]) - assert "max_input_tokens" not in params - assert params["model"] == "gpt-4o-mini" - assert params["input_cost_per_token"] == 0.000001 - assert info["max_input_tokens"] == 4096 - - def test_omitted_key_is_untouched_by_a_null_elsewhere(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"mode": None, "supports_vision": False}} - ), - ) - - info = json.loads(result["model_info"]) - assert "mode" not in info - assert info["supports_vision"] is False - assert info["max_input_tokens"] == 4096 - - @pytest.mark.parametrize("field", sorted(_PROTECTED_MODEL_INFO_VALUES)) - def test_null_on_protected_key_is_ignored(self, field): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate({"model_info": {field: None}}), - ) - - info = json.loads(result["model_info"]) - assert info[field] == _PROTECTED_MODEL_INFO_VALUES[field] - assert info["max_input_tokens"] == 4096 - - def test_echoing_the_read_back_blob_preserves_every_stored_key(self): - """The Admin UI edit form submits the whole /model/info row back, and that read reports - every key the deployment never stored as an explicit null. Those nulls have to stay - no-ops: a write drops None before storing, so a null in the echoed blob always names a - key the stored row does not carry. - """ - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - db_model = _build_db_model_with_pinned_model_info() - echoed = { - "id": "dep-pinned-0", - "max_input_tokens": 4096, - "mode": "chat", - "supports_vision": True, - "input_cost_per_token": 0.000001, - "team_id": "team-keep-me", - "base_model": None, - "tier": None, - "max_output_tokens": None, - "supports_function_calling": None, - "cache_read_input_token_cost": None, - } - - result = update_db_model( - db_model=db_model, - updated_patch=updateDeployment.model_validate({"model_info": echoed}), - ) - - info = json.loads(result["model_info"]) - assert info["max_input_tokens"] == 4096 - assert info["mode"] == "chat" - assert info["supports_vision"] is True - assert info["input_cost_per_token"] == 0.000001 - assert info["team_id"] == "team-keep-me" - for never_stored in ("base_model", "tier", "max_output_tokens", "supports_function_calling"): - assert never_stored not in info - - def test_null_on_pricing_key_still_clears_both_blobs(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"input_cost_per_token": None}} - ), - ) - - params = json.loads(result["litellm_params"]) - info = json.loads(result["model_info"]) - assert "input_cost_per_token" not in params - assert "input_cost_per_token" not in info - assert params["max_input_tokens"] == 4096 - assert info["max_input_tokens"] == 4096 - - class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index d22ec60e61a..30b265905f3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,198 +220,6 @@ def test_should_store_full_pricing_under_deployment_model_id(): assert entry["output_cost_per_token"] == 0.0 -def test_should_drop_a_price_the_deployment_no_longer_carries(): - """Re-registering a deployment must replace its model_id entry, not merge onto it. - - A merge left the old rate in the cost map after an operator cleared the override, so - the deployment kept billing at a price its config no longer had. - """ - backend_model = "vertex_ai/gemini-2.5-flash" - model_id = "deployment-cleared-price" - original = {model_id: litellm.model_cost.get(model_id)} - - try: - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"input_cost_per_token": 0.005, "output_cost_per_token": 0.01}, - model=backend_model, - custom_llm_provider="vertex_ai", - ) - assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 - - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"mode": "chat"}, - model=backend_model, - custom_llm_provider="vertex_ai", - ) - - entry = litellm.model_cost[model_id] - assert entry.get("input_cost_per_token") != 0.005, ( - "the cleared override survived re-registration, so the deployment still bills at it" - ) - assert entry.get("output_cost_per_token") != 0.01 - finally: - _restore_model_cost_entries(original) - - -def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it(): - """Deployments are keyed into the same cost map as the built-in catalog, so a deployment - whose id happens to name a real model must not evict that model's entry. - - Stripping it would take the pricing and capability flags every other deployment of that - model reads, process-wide, until the next price-map reload. Registering twice, because - the first registration is what would mark the entry as this deployment's own. - """ - colliding_id = "gpt-4o" - original = {colliding_id: litellm.model_cost.get(colliding_id)} - builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] - - try: - for _ in range(2): - Router._register_deployment_in_model_cost( - model_id=colliding_id, - model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - entry = litellm.model_cost[colliding_id] - assert entry["max_tokens"] == builtin_max_tokens, ( - "registering a deployment under a catalog model's name wiped that model's context window" - ) - assert entry["litellm_provider"] == "openai" - assert entry["supports_vision"] is True - finally: - _restore_model_cost_entries(original) - - -def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): - """A deployment may carry `litellm_provider` in its own model_info, which must not be - read as "this is a catalog entry" and stop the stale price from being dropped.""" - model_id = "deployment-provider-tagged" - original = {model_id: litellm.model_cost.get(model_id)} - - try: - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 - - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( - "a deployment that declares its provider kept billing at the price it no longer carries" - ) - finally: - _restore_model_cost_entries(original) - - -def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): - """Deleting a deployment releases its claim on the shared cost-map key. - - Held forever, a later catalog refresh that starts publishing a model under that same - name would be treated as the deleted deployment's own entry and evicted. - """ - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-to-delete" - original = {model_id: litellm.model_cost.get(model_id)} - router = Router( - model_list=[ - { - "model_name": "to-delete", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - ] - ) - - try: - assert model_id in _DEPLOYMENT_COST_MAP_KEYS - - assert router.delete_deployment(id=model_id) is not None - - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS, ( - "a deleted deployment kept its claim on the shared cost-map key" - ) - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - -def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): - """Two live routers can serve the same deployment id, and the claim is process-wide. - - Releasing it when only one of them drops the deployment would put the survivor back on - merging, so the price it just cleared would keep billing. - """ - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-served-twice" - original = {model_id: litellm.model_cost.get(model_id)} - entry = { - "model_name": "served-twice", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - first = Router(model_list=[entry]) - second = Router(model_list=[entry]) - - try: - assert model_id in _DEPLOYMENT_COST_MAP_KEYS - - assert first.delete_deployment(id=model_id) is not None - - assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( - "the claim was released while another router still served the deployment" - ) - - assert second.delete_deployment(id=model_id) is not None - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - -def test_should_keep_the_cost_map_key_while_a_dynamically_built_router_serves_it(): - """A router built with no model_list still serves whatever add_deployment gives it, so it - counts when deciding whether the shared cost-map claim can be released.""" - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-added-dynamically" - original = {model_id: litellm.model_cost.get(model_id)} - entry = { - "model_name": "added-dynamically", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - configured = Router(model_list=[entry]) - dynamic = Router() - dynamic.add_deployment(deployment=Deployment(**entry)) - - try: - assert configured.delete_deployment(id=model_id) is not None - - assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( - "the claim was released while a dynamically built router still served the deployment" - ) - - assert dynamic.delete_deployment(id=model_id) is not None - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 93aaf3ca58c..83b0d58f2b2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9064,9 +9064,8 @@ export interface paths { * Patch Model * @description PATCH Endpoint for partial model updates. * - * JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the - * body is unchanged, a key sent as null is removed from the stored row, and a key sent with a - * value is set (identity and ownership keys such as `id` and `team_id` ignore a null). + * Only updates the fields specified in the request while preserving other existing values. + * Follows proper PATCH semantics by only modifying provided fields. * * Args: * model_id: The ID of the model to update From 9f21ae395aa298ab61331fb74e85d4769aae05e5 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 13:16:04 +0000 Subject: [PATCH 093/157] fix(registry): correct eu Claude 3.5 Haiku Bedrock pricing, add Nova v1 tool_choice, Azure gpt-5.5 snapshot retirement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 65 ++++++++++++------- model_prices_and_context_window.json | 65 ++++++++++++------- .../test_ai_policy_suggester.py | 2 +- tests/test_litellm/test_utils.py | 40 ++++++++---- 4 files changed, 111 insertions(+), 61 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 43def22ba58..d1f1a2d2b48 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -537,7 +538,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +553,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -2876,7 +2879,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2888,7 +2892,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2902,7 +2907,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -8064,7 +8070,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8108,7 +8115,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8152,7 +8160,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -12317,7 +12326,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12496,7 +12506,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12508,7 +12519,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12522,7 +12534,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -20990,7 +21003,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -21002,7 +21016,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -21017,24 +21032,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -43795,7 +43811,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43807,7 +43824,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43836,7 +43854,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 43def22ba58..d1f1a2d2b48 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -537,7 +538,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +553,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -2876,7 +2879,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2888,7 +2892,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2902,7 +2907,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -8064,7 +8070,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8108,7 +8115,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8152,7 +8160,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -12317,7 +12326,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12496,7 +12506,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12508,7 +12519,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12522,7 +12534,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -20990,7 +21003,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -21002,7 +21016,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -21017,24 +21032,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -43795,7 +43811,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43807,7 +43824,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43836,7 +43854,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index 93dc429168f..bb71d67f24e 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -265,7 +265,7 @@ class TestSuggesterRejectsModelsWithoutToolCalling: def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): supported_params = litellm.get_supported_openai_params( - model="amazon.nova-pro-v1:0", + model="meta.llama4-scout-17b-instruct-v1:0", custom_llm_provider="bedrock", ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d10e7634720..ec3a141a490 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1409,23 +1409,35 @@ def test_supports_tool_choice_simple_tests(): is True ) - assert ( - litellm.utils.supports_tool_choice(model="us.amazon.nova-micro-v1:0") is False - ) - assert ( - litellm.utils.supports_tool_choice(model="bedrock/us.amazon.nova-micro-v1:0") - is False - ) - assert ( - litellm.utils.supports_tool_choice( - model="us.amazon.nova-micro-v1:0", custom_llm_provider="bedrock_converse" - ) - is False - ) - assert litellm.utils.supports_tool_choice(model="perplexity/sonar") is False +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model", + [ + "amazon.nova-lite-v1:0", + "amazon.nova-micro-v1:0", + "amazon.nova-pro-v1:0", + "apac.amazon.nova-lite-v1:0", + "apac.amazon.nova-micro-v1:0", + "apac.amazon.nova-pro-v1:0", + "bedrock/us-gov-east-1/amazon.nova-pro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-lite-v1:0", + "bedrock/us-gov-west-1/amazon.nova-micro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-pro-v1:0", + "eu.amazon.nova-lite-v1:0", + "eu.amazon.nova-micro-v1:0", + "eu.amazon.nova-pro-v1:0", + "us.amazon.nova-lite-v1:0", + "us.amazon.nova-micro-v1:0", + "us.amazon.nova-pro-v1:0", + ], +) +def test_amazon_nova_v1_understanding_models_support_tool_choice(model: str) -> None: + assert litellm.utils.supports_tool_choice(model=model) is True + + def test_check_provider_match(): """ Test the _check_provider_match function for various provider scenarios From b067729082783e02c13a7b757360caa09d1297be Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 13:51:40 +0000 Subject: [PATCH 094/157] fix(registry): add xAI Imagine Video 720p per second rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++++ model_prices_and_context_window.json | 4 ++++ model_prices_and_context_window.schema.json | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d1f1a2d2b48..dc720e2629a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59849,6 +59849,7 @@ "mode": "video_generation", "output_cost_per_second": 0.05, "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, "source": "https://docs.x.ai/docs/models/grok-imagine-video", "supported_modalities": [ "text", @@ -59866,6 +59867,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", @@ -59883,6 +59885,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", @@ -59900,6 +59903,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d1f1a2d2b48..dc720e2629a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59849,6 +59849,7 @@ "mode": "video_generation", "output_cost_per_second": 0.05, "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, "source": "https://docs.x.ai/docs/models/grok-imagine-video", "supported_modalities": [ "text", @@ -59866,6 +59867,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", @@ -59883,6 +59885,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", @@ -59900,6 +59903,7 @@ "output_cost_per_second": 0.08, "output_cost_per_second_1080p": 0.25, "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", "supported_modalities": [ "text", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 47a1934a703..7ed1e7e568b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -478,6 +478,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_720p": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, From 2881b8cd4553f1dab28e7220e7f5a1f30e960aed Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 14:09:25 +0000 Subject: [PATCH 095/157] fix(cost): carry output_cost_per_second_720p through model info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/router.py | 1 + litellm/types/utils.py | 2 ++ litellm/utils.py | 1 + tests/test_litellm/test_utils.py | 2 ++ tests/test_litellm/test_video_generation.py | 26 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +++ 6 files changed, 36 insertions(+) diff --git a/litellm/types/router.py b/litellm/types/router.py index 5c9eab30f3d..6b707a544a2 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -525,6 +525,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): input_cost_per_second: float | None output_cost_per_second: float | None output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_1080p: float | None output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d62f00f3676..ab0cc5f959c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -318,6 +318,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_credit: float | None # for OCR models priced by credit @@ -3522,6 +3523,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None output_cost_per_second_480p: float | None = None + output_cost_per_second_720p: float | None = None output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 36b48d3b8d8..d0ff4917616 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5913,6 +5913,7 @@ def _get_model_info_helper( output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), + output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index ec3a141a490..357c0d1c363 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -814,6 +814,7 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_second", "output_cost_per_second", "output_cost_per_second_480p", + "output_cost_per_second_720p", "output_cost_per_second_1080p", "output_cost_per_second_4k", "input_cost_per_query", @@ -1037,6 +1038,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, + "output_cost_per_second_720p": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 2a60ff9c4b5..f3cd4618078 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,32 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_xai_imagine_video_720p_tier_from_cost_map(self, monkeypatch): + """720p xAI Imagine Video requests bill the published 720p rate, not the 480p base rate.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, resolution: str, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = {"duration_seconds": duration, "video_resolution": resolution} + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider="xai", + ) + + assert abs(cost_for("xai/grok-imagine-video", "720p", 10.0) - 0.7) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "720p", 10.0) - 1.4) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "480p", 10.0) - 0.8) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "1080p", 10.0) - 2.5) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" from litellm.cost_calculator import completion_cost diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..ce50978602f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29629,6 +29629,8 @@ export interface components { output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; + /** Output Cost Per Second 720P */ + output_cost_per_second_720p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ @@ -39820,6 +39822,8 @@ export interface components { output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; + /** Output Cost Per Second 720P */ + output_cost_per_second_720p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ From a1588c260232dd6ed22605fdd4db24983bb1fcee Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 08:05:15 -0700 Subject: [PATCH 096/157] fix(mcp): preserve stream failures across transports --- litellm/experimental_mcp_client/client.py | 7 +- .../mcp_server/rest_endpoints.py | 25 +- .../test_mcp_client.py | 240 ++++++++++++++++++ .../mcp_server/test_rest_endpoints.py | 48 ++++ 4 files changed, 313 insertions(+), 7 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index cabcc6c03ba..5cdc4efdfef 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -451,7 +451,7 @@ class MCPClient: async def receive_message( message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, ) -> None: - if not isinstance(message, ValueError): + if not isinstance(message, (ValueError, httpx.RequestError, OSError)): return if not stream_error.done(): stream_error.set_result(message) @@ -472,7 +472,7 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=timedelta(seconds=self.timeout), - message_handler=receive_message if self.transport_type == MCPTransport.http else None, + message_handler=receive_message, **session_kwargs, ) session: Final = await session_ctx.__aenter__() @@ -525,8 +525,7 @@ class MCPClient: read_timeout: Final = _as_read_timeout(e) if read_timeout is not None: verbose_logger.warning( - "MCP client timed out after %ss waiting for %s to answer; the server accepted the " - "request and ended its response stream without a JSON-RPC reply", + "MCP client timed out after %ss waiting for a valid MCP response from %s", self.timeout, self.server_url or "stdio", ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index c0ffeaeac62..63badc6f703 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -114,7 +114,8 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout return str(exc.detail) if isinstance(exc, TimeoutError): return ( - f"Failed to connect to MCP server: no response from {_redact_mcp_resource_url(url) or 'the server'} " + "Failed to connect to MCP server: no valid MCP response received from " + f"{_redact_mcp_resource_url(url) or 'the server'} " f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) @@ -131,6 +132,11 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." + if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)): + return ( + "Failed to connect to MCP server: the connection was interrupted. " + "Check the server and network connection, then retry." + ) if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"): return ( "Failed to connect to MCP server: the endpoint returned an unsupported content type. " @@ -142,6 +148,11 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout "Check the MCP endpoint URL and the server's protocol implementation." ) if MCP_AVAILABLE and isinstance(exc, McpError): + if exc.error.code == -32000 and exc.error.message == "Connection closed": + return ( + "Failed to connect to MCP server: the connection was closed before the request completed. " + "Check that the server stays running and returns a complete MCP response, then retry." + ) if exc.error.code == 32600 and exc.error.message == "Session terminated": return ( "Failed to connect to MCP server: the MCP session was terminated. " @@ -159,7 +170,7 @@ if MCP_AVAILABLE: from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool - from litellm.experimental_mcp_client.client import MCPClient + from litellm.experimental_mcp_client.client import MCPClient, _as_read_timeout from litellm.llms.litellm_proxy.skills.skill_search import ( DEFAULT_SKILL_SEARCH_TOP_K, ) @@ -1396,10 +1407,18 @@ if MCP_AVAILABLE: except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise except BaseException as e: + effective_timeout: Final = ( + min(request.timeout or MCP_CLIENT_TIMEOUT, timeout_seconds) + if any( + isinstance(cause, McpError) and _as_read_timeout(cause) is not None + for cause in iter_exception_tree(e) + ) + else timeout_seconds + ) return { "status": "error", "error": True, - "message": _connection_error_message(e, request.url, timeout_seconds), + "message": _connection_error_message(e, request.url, effective_timeout), } async def _preview_openapi_tools(spec_path: str) -> dict: diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 90292b9b162..a1e781d17b7 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -3,6 +3,7 @@ import base64 import json import os import sys +from collections.abc import AsyncIterator from importlib import metadata from pathlib import Path from typing import Final @@ -18,6 +19,7 @@ from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( LATEST_PROTOCOL_VERSION, + CallToolResult, ErrorData, Implementation, InitializeResult, @@ -36,6 +38,7 @@ from litellm.experimental_mcp_client.client import ( MCPClient, _as_read_timeout, _first_non_cancelled_cause, + _TransportContext, missing_streamable_http_client_error, strip_auth_scheme, ) @@ -1296,6 +1299,7 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: [ ("text/html", b"secret-page", ValueError), ("application/json", b"secret-invalid-json", ValidationError), + ("application/json", b"", ValidationError), ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), ], @@ -1446,3 +1450,239 @@ async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() message: Final = _connection_error_message(caught.value, client.server_url, 30) assert "invalid MCP response" in message assert "secret" not in message + + +class _DiagnosticSSEStream(httpx.AsyncByteStream): + def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: + self.messages = messages + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + message: Final = await self.messages.get() + if message is None: + return + if isinstance(message, Exception): + raise message + yield b"event: message\ndata: " + message + b"\n\n" + + +_DIAGNOSTIC_STDIO_SERVER: Final = """ +import json, sys +mode, failure_method = sys.argv[1:] +for line in sys.stdin: + request = json.loads(line) + if "method" not in request or "id" not in request: + continue + if request["method"] == failure_method: + if mode == "bad-json": + print("secret-invalid-json", flush=True) + continue + if mode == "closed": + sys.exit(0) + if mode == "silent": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) + continue + if request["method"] == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} + elif request["method"] == "tools/list": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) + result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + else: + result = {"content": [{"type": "text", "text": "pong"}], "isError": False} + print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) +""" + + +def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: + from mcp import StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + + if transport == MCPTransport.stdio: + return stdio_client( + StdioServerParameters( + command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] + ) + ) + messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, headers={"Content-Type": "text/event-stream"}, stream=_DiagnosticSSEStream(messages) + ) + payload: Final = json.loads(request.content) + if "method" not in payload or "id" not in payload: + return httpx.Response(202) + if payload["method"] == failure_method and mode != "ok": + if mode == "bad-json": + await messages.put(b"secret-invalid-json") + elif mode == "io-error": + await messages.put(httpx.ReadError("secret-read-error")) + elif mode == "closed": + await messages.put(None) + elif mode == "silent": + await messages.put( + b'{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Waiting"}}' + ) + return httpx.Response(202) + if payload["method"] == "tools/list": + for message in ( + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + }, + {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, + {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, + ): + await messages.put(json.dumps(message).encode()) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"tools": {}, "logging": {}}, + "serverInfo": {"name": "diagnostic", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + if payload["method"] == "tools/list" + else {"content": [{"type": "text", "text": "pong"}], "isError": False} + ) + await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) + return httpx.Response(202) + + def factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) + + return sse_client("https://example.com/sse", httpx_client_factory=factory) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) +async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) + with pytest.raises(ValidationError): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_sse_read_failure_is_preserved() -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) + with pytest.raises(httpx.ReadError, match="secret-read-error"): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(MCPTransport.sse, "io-error", "tools/list"), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: + from mcp import ClientSession + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + logging_callback: Final = AsyncMock() + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=0.2, logging_callback=logging_callback + ) + + async def operation(session: ClientSession) -> CallToolResult: + tools: Final = await session.list_tools() + assert [tool.name for tool in tools.tools] == ["ping"] + return await session.call_tool("ping", {}) + + pending: Final = client._execute_session_operation(_diagnostic_transport(transport, mode, "tools/list"), operation) + if mode == "ok": + result: Final = await asyncio.wait_for(pending, timeout=3) + assert result.isError is False + assert result.content[0].text == "pong" + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + else: + with pytest.raises(McpError) as caught: + await asyncio.wait_for(pending, timeout=3) + if mode == "closed": + assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) + else: + assert isinstance(_as_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: + ready: Final = asyncio.Event() + + async def on_log(message: LoggingMessageNotificationParams) -> None: + if message.data == "Waiting": + ready.set() + + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log + ) + task: Final = asyncio.create_task( + client._execute_session_operation( + _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() + ) + ) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +class _InterruptedHTTPBody(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'{"jsonrpc":' + raise httpx.RemoteProtocolError("secret-incomplete-response") + + +@pytest.mark.asyncio +async def test_interrupted_http_response_preserves_the_transport_failure() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(httpx.RemoteProtocolError, match="secret-incomplete-response"): + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=0.2) + with pytest.raises(McpError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + assert isinstance(_as_read_timeout(caught.value), TimeoutError) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 73734131130..101b02e6ffd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3427,6 +3427,54 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message + @pytest.mark.parametrize( + "error_type", [httpx.ReadError, httpx.WriteError, httpx.RemoteProtocolError, ConnectionResetError] + ) + def test_interrupted_connection_message_is_safe(self, error_type: type[Exception]) -> None: + message: Final = rest_endpoints._connection_error_message( + error_type("secret-transport-detail"), "https://example.com/?token=secret-query", 30 + ) + assert "connection was interrupted" in message + assert "secret" not in message + + def test_closed_connection_explains_incomplete_request(self) -> None: + from mcp import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=-32000, message="Connection closed", data="secret-data")), None, 30 + ) + assert "connection was closed before the request completed" in message + assert "secret" not in message + + def test_timeout_does_not_claim_the_server_sent_nothing(self) -> None: + message: Final = rest_endpoints._connection_error_message(TimeoutError(), None, 30) + assert "no valid MCP response received" in message + + @pytest.mark.asyncio + @pytest.mark.parametrize("sdk_timeout", [True, False]) + async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool) -> None: + from mcp import McpError + from mcp.types import ErrorData + + async def operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + try: + raise TimeoutError("secret-timeout") + except TimeoutError as elapsed: + if not sdk_timeout: + raise + try: + raise McpError(ErrorData(code=408, message="secret-sdk-timeout")) from elapsed + except McpError as sdk_error: + raise TimeoutError() from sdk_error + + payload: Final = NewMCPServerRequest( + server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=1 + ) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30) + assert ("within 1s" if sdk_timeout else "within 30s") in result["message"] + assert "secret" not in result["message"] + def test_unknown_error_falls_back_to_generic(self): message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message From 1fc1aaabdaa3542cd2d410103ee01ebc94025961 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 16:00:35 +0000 Subject: [PATCH 097/157] test(cli): drop lite --version subprocess regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/test_global_options.py | 56 ------------------- 1 file changed, 56 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 9ed9f91ab9f..b73d1acc6e3 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,10 +1,6 @@ # stdlib imports import json import os -import shutil -import subprocess -import threading -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest.mock import Mock, patch @@ -37,58 +33,6 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output -def test_lite_version_does_not_fetch_model_cost_map(tmp_path: Path) -> None: - lite_path = shutil.which("lite") - if lite_path is None: - pytest.skip("lite executable is unavailable") - - request_log: Path = tmp_path / "requests.log" - - class _CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - with request_log.open("a", encoding="utf-8") as log_file: - log_file.write(f"{self.path}\n") - body = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, format: str, *args: object) -> None: - return - - server = ThreadingHTTPServer(("127.0.0.1", 0), _CostMapHandler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - source_root = str(Path(__file__).resolve().parents[5]) - env["PYTHONPATH"] = os.pathsep.join(filter(None, (source_root, env.get("PYTHONPATH")))) - env.update( - { - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": "http://127.0.0.1:9", - } - ) - result = subprocess.run( - [lite_path, "--version"], - capture_output=True, - text=True, - timeout=120, - env=env, - ) - finally: - server.shutdown() - server.server_close() - thread.join(timeout=10) - - assert result.returncode == 0 - assert "LiteLLM Proxy CLI Version" in result.stdout - request_count: int = request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 - assert request_count == 0 - - def test_cli_source_is_ascii_only(): """Non-ASCII output (emoji, box-drawing chars) raises UnicodeEncodeError on legacy Windows consoles (cp1252), so the whole CLI package must stay ASCII-only.""" From 9b6c7c8bf0e2e0e3218caa8a18b29b17bb8bf566 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:17:14 -0700 Subject: [PATCH 098/157] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/_experimental/mcp_server/rest_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 63badc6f703..31252479445 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1408,7 +1408,7 @@ if MCP_AVAILABLE: raise except BaseException as e: effective_timeout: Final = ( - min(request.timeout or MCP_CLIENT_TIMEOUT, timeout_seconds) + min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) if any( isinstance(cause, McpError) and _as_read_timeout(cause) is not None for cause in iter_exception_tree(e) From 6dbaf43ba2951e403fd1efe6920eaf6c30550c42 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 16:42:07 +0000 Subject: [PATCH 099/157] fix(registry): drop OpenAI shutdown date from shared computer-use-preview entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 - model_prices_and_context_window.json | 1 - 2 files changed, 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dc720e2629a..cdd9fc7b6d4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14560,7 +14560,6 @@ "supports_tool_choice": true }, "computer-use-preview": { - "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, "litellm_provider": "azure", "max_input_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dc720e2629a..cdd9fc7b6d4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14560,7 +14560,6 @@ "supports_tool_choice": true }, "computer-use-preview": { - "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, "litellm_provider": "azure", "max_input_tokens": 8192, From dbf94902294cf5fa4fe5af6681fcd83ffd743231 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:44:14 -0700 Subject: [PATCH 100/157] fix(mcp): expose shared SDK timeout normalization --- litellm/experimental_mcp_client/client.py | 6 +++--- .../_experimental/mcp_server/rest_endpoints.py | 4 ++-- .../experimental_mcp_client/test_mcp_client.py | 18 +++++++++--------- .../mcp_server/test_rest_endpoints.py | 7 ++++--- 4 files changed, 18 insertions(+), 17 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 5cdc4efdfef..8c7f4557992 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -150,8 +150,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT) otherwise carries JSON-RPC error codes.""" -def _as_read_timeout(exc: BaseException) -> TimeoutError | None: - """The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``. +def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: + """Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``. The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error @@ -522,7 +522,7 @@ class MCPClient: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception as e: - read_timeout: Final = _as_read_timeout(e) + read_timeout: Final = as_mcp_read_timeout(e) if read_timeout is not None: verbose_logger.warning( "MCP client timed out after %ss waiting for a valid MCP response from %s", diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 31252479445..5b74caee3a2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -170,7 +170,7 @@ if MCP_AVAILABLE: from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool - from litellm.experimental_mcp_client.client import MCPClient, _as_read_timeout + from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout from litellm.llms.litellm_proxy.skills.skill_search import ( DEFAULT_SKILL_SEARCH_TOP_K, ) @@ -1410,7 +1410,7 @@ if MCP_AVAILABLE: effective_timeout: Final = ( min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) if any( - isinstance(cause, McpError) and _as_read_timeout(cause) is not None + isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None for cause in iter_exception_tree(e) ) else timeout_seconds diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index a1e781d17b7..b07e5876e8b 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -36,9 +36,9 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import ( MCP_STREAMABLE_HTTP_REQUIREMENT, MCPClient, - _as_read_timeout, _first_non_cancelled_cause, _TransportContext, + as_mcp_read_timeout, missing_streamable_http_client_error, strip_auth_scheme, ) @@ -867,25 +867,25 @@ def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> McpErr return raised -def test_as_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): +def test_as_mcp_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): """Neither signal alone is enough. The code alone cannot separate the SDK's own timeout from an upstream JSON-RPC error that happens to use 408, and the context chain alone cannot separate it from any other relayed error that surfaces while a timeout is being handled, so both must hold. """ timeout_code = int(httpx.codes.REQUEST_TIMEOUT) - translated = _as_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) + translated = as_mcp_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) assert isinstance(translated, TimeoutError) assert str(translated) == "Timed out while waiting" relayed_408 = McpError(ErrorData(code=timeout_code, message="upstream said 408")) - assert _as_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" + assert as_mcp_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" relayed_other = _raise_mcp_error_while_handling_a_timeout(-32603, "upstream internal error") - assert _as_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" + assert as_mcp_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" - assert _as_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None - assert _as_read_timeout(RuntimeError("not an McpError")) is None + assert as_mcp_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None + assert as_mcp_read_timeout(RuntimeError("not an McpError")) is None @pytest.mark.asyncio @@ -1619,7 +1619,7 @@ async def test_transport_completion_and_normal_messages(transport: MCPTransport, if mode == "closed": assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) else: - assert isinstance(_as_read_timeout(caught.value), TimeoutError) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) @pytest.mark.asyncio @@ -1685,4 +1685,4 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N ), timeout=3, ) - assert isinstance(_as_read_timeout(caught.value), TimeoutError) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 101b02e6ffd..3487f634251 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3453,7 +3453,8 @@ class TestConnectionErrorMessage: @pytest.mark.asyncio @pytest.mark.parametrize("sdk_timeout", [True, False]) - async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool) -> None: + @pytest.mark.parametrize("read_timeout", [0, 1]) + async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool, read_timeout: int) -> None: from mcp import McpError from mcp.types import ErrorData @@ -3469,10 +3470,10 @@ class TestConnectionErrorMessage: raise TimeoutError() from sdk_error payload: Final = NewMCPServerRequest( - server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=1 + server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=read_timeout ) result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30) - assert ("within 1s" if sdk_timeout else "within 30s") in result["message"] + assert (f"within {read_timeout}s" if sdk_timeout else "within 30s") in result["message"] assert "secret" not in result["message"] def test_unknown_error_falls_back_to_generic(self): From 1183b2abc6645a92f9a852daa7843662b6fa6f20 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 9 Sep 2026 09:46:37 -0700 Subject: [PATCH 101/157] fix(integrations): pass original request object to post-call guardrail hooks (#40414) --- litellm/integrations/custom_guardrail.py | 30 ++++--- .../integrations/test_custom_guardrail.py | 90 ++++++++++++++++++- 2 files changed, 106 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 37d6a7e793d..77bf4820a1a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -850,20 +850,24 @@ class CustomGuardrail(CustomLogger): if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: return None - # CHECK IF GUARDRAIL REJECTS THE REQUEST target: Final = self._deployment_hook_target() - hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data - result: Final = await target.async_post_call_success_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id=request_data.get("user_api_key_user_id"), - team_id=request_data.get("user_api_key_team_id"), - end_user_id=request_data.get("user_api_key_end_user_id"), - api_key=request_data.get("user_api_key_hash"), - request_route=request_data.get("user_api_key_request_route"), - ), - data=hook_request_data, - response=response, - ) + try: + if target is not self: + request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key + result: Final = await target.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id=request_data.get("user_api_key_user_id"), + team_id=request_data.get("user_api_key_team_id"), + end_user_id=request_data.get("user_api_key_end_user_id"), + api_key=request_data.get("user_api_key_hash"), + request_route=request_data.get("user_api_key_request_route"), + ), + data=request_data, + response=response, + ) + finally: + if target is not self: + request_data.pop("guardrail_to_apply", None) if not self._is_valid_response_type(result): return None diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index cd8d609cf71..ddc8439a83a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1842,12 +1842,14 @@ class _ApplyStyleGuardrail(CustomGuardrail): self.block = block self.apply_called = False self.seen_texts = None + self.seen_request_data = None async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): from fastapi import HTTPException self.apply_called = True self.seen_texts = inputs.get("texts") + self.seen_request_data = request_data if self.block: raise HTTPException(status_code=400, detail={"error": "Violated moderation policy"}) return inputs @@ -2646,6 +2648,91 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: which starved every later callback in litellm.callbacks (notably the lazily-appended VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + @pytest.mark.asyncio + async def test_apply_guardrail_retains_request_identity(self) -> None: + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = _ApplyStyleGuardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="review me"))]) + + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=CallTypes.acompletion + ) + + assert guardrail.seen_request_data is request_data + assert guardrail.seen_texts == ["review me"] + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", (None, CallTypes.acompletion)) + async def test_apply_guardrail_masks_response_and_records_metadata(self, call_type: CallTypes | None) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="response-filter", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="secret", action=ContentFilterAction.MASK)], + ) + request_data: Final = {"guardrails": ["response-filter"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="a secret"))]) + + result: Final = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=call_type + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == f"a {guardrail.keyword_redaction_tag}" + entries: Final = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "response-filter" + assert entries[0]["guardrail_mode"] == "post_call" + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("error_type", (None, RuntimeError, asyncio.CancelledError)) + async def test_dispatch_cleans_up_request_on_every_exit(self, error_type: type[BaseException] | None) -> None: + from contextlib import nullcontext + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import LLMResponseTypes, ModelResponse + + error: Final = error_type("dispatch interrupted") if error_type is not None else None + + class Dispatch(CustomLogger): + request_data: dict[str, object] | None = None + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> LLMResponseTypes: + self.request_data = data + if error is not None: + raise error + return response + + dispatch: Final = Dispatch() + + class Guardrail(_ApplyStyleGuardrail): + def _deployment_hook_target(self) -> CustomLogger: + return dispatch + + guardrail: Final = Guardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + with pytest.raises(error_type) if error_type is not None else nullcontext(): + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=ModelResponse(), call_type=CallTypes.acompletion + ) + + assert dispatch.request_data is request_data + assert "guardrail_to_apply" not in request_data + @pytest.mark.asyncio async def test_returns_none_when_request_has_no_guardrails(self): from litellm.types.utils import ModelResponse @@ -2740,4 +2827,5 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: assert result is response assert response.choices[0].message.content == "filtered response" - assert request_data == {"guardrails": ["test-guardrail"]} + assert "guardrail_to_apply" not in request_data + assert len(_guardrail_entries(request_data)) == 1 From 36f3ca95d8bdb6b03ab4632982cdc7c0cef99c6f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 09:57:38 -0700 Subject: [PATCH 102/157] fix(proxy): report /cost/estimate rates from the call that billed them The estimate looked the reported per-token rates up a second time, with the provider this endpoint resolved rather than the one completion_cost infers. The provider decides whether a token tier threshold is inclusive, so an unrouted xai model sitting exactly on 200k billed at the tier rate and reported the base rate, half of it. completion_cost now hands back the rates its own lines were billed at, and the endpoint reports those. Claude-Session: https://claude.ai/code/session_01RLKy5DMi3XCBUJ37WzfNi1 --- litellm/cost_calculator.py | 6 ++ litellm/litellm_core_utils/litellm_logging.py | 5 ++ .../litellm_core_utils/llm_cost_calc/utils.py | 60 ++++++++++--------- .../cost_tracking_settings.py | 29 ++++----- .../llm_cost_calc/test_llm_cost_calc_utils.py | 52 +++++++++++++--- .../test_cost_tracking_settings.py | 39 ++++++++++++ tests/test_litellm/test_cost_calculator.py | 56 +++++++++++++++++ 7 files changed, 195 insertions(+), 52 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a1181ee4124..b4be023540f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import TranscriptionUsageObjectTransformation, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, _generic_cost_per_character, _get_regional_uplift_multiplier, @@ -1122,6 +1123,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: BilledTokenRates | None = None, ) -> None: """ Helper function to store cost breakdown in the logging object. @@ -1166,6 +1168,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=billed_token_rates, ) except Exception as breakdown_error: @@ -1735,6 +1738,7 @@ def completion_cost( _reasoning_cost: float | None = None _cache_read_cost: float | None = None _cache_creation_cost: float | None = None + _billed_token_rates: BilledTokenRates | None = None if cost_per_token_usage_object is not None and model: _breakdown_provider: str | None = ( custom_llm_provider if isinstance(custom_llm_provider, str) else None @@ -1751,6 +1755,7 @@ def completion_cost( _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost _cache_creation_cost = _token_type_breakdown.cache_creation_cost + _billed_token_rates = _token_type_breakdown.rates _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar, @@ -1770,6 +1775,7 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=_billed_token_rates, ) return _final_cost diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ca2cca5360f..a9f3c0e9325 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -201,6 +201,7 @@ if TYPE_CHECKING: from mcp.types import EmbeddedResource, ImageContent, TextContent from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( @@ -581,6 +582,7 @@ class Logging(LiteLLMLoggingBaseClass): # Initialize cost breakdown field self.cost_breakdown: CostBreakdown | None = None + self.billed_token_rates: BilledTokenRates | None = None # Init Caching related details self.caching_details: CachingDetails | None = None @@ -1585,6 +1587,7 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: "BilledTokenRates | None" = None, ) -> None: """ Helper method to store cost breakdown in the logging object. @@ -1604,8 +1607,10 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: Tier the costs above were priced on, already resolved data_residency: Region uplift the costs above were priced on, already resolved vertex_location: Vertex AI location the costs above were priced on, already resolved + billed_token_rates: Per-token rates the costs above were billed at, already resolved """ + self.billed_token_rates = billed_token_rates self.cost_breakdown = CostBreakdown( input_cost=input_cost, output_cost=output_cost, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e2168528e6b..42f88bbb425 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1302,34 +1302,6 @@ def _coerce_token_count(value: object) -> int: return value if isinstance(value, int) and value > 0 else 0 -@dataclass(frozen=True, slots=True) -class TokenTypeCostBreakdown: - reasoning_cost: float - cache_read_cost: float - cache_creation_cost: float - - -def _reasoning_token_count(usage: Usage) -> int: - parsed: Final = ( - parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 - ) - return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - - -def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: - """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details - first, then the private top-level counters the Usage constructor mirrors cache tokens onto for - providers/callers that bypass the details.""" - parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None - parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 - parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 - return ( - parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), - parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), - parsed["cache_creation_token_details"] if parsed is not None else None, - ) - - @dataclass(frozen=True, slots=True) class BilledTokenRates: """Per-token rates one request's usage bills at, after token tiers, off-peak windows and the @@ -1355,6 +1327,37 @@ class BilledTokenRates: ) +@dataclass(frozen=True, slots=True) +class TokenTypeCostBreakdown: + reasoning_cost: float + cache_read_cost: float + cache_creation_cost: float + rates: BilledTokenRates | None = None + """Rates these lines were billed at, so a caller reporting both cannot resolve them a second, + differently-argued way. None when the model's pricing could not be resolved.""" + + +def _reasoning_token_count(usage: Usage) -> int: + parsed: Final = ( + parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 + ) + return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) + + +def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details + first, then the private top-level counters the Usage constructor mirrors cache tokens onto for + providers/callers that bypass the details.""" + parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None + parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 + parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 + return ( + parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), + parsed["cache_creation_token_details"] if parsed is not None else None, + ) + + def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates: """Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" @@ -1497,6 +1500,7 @@ def get_token_type_cost_breakdown( reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, cache_creation_cost=cache_creation_cost, + rates=rates, ) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 1faa66584d5..493f75008c3 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -21,7 +21,6 @@ import litellm from litellm._internal_context import current_billing_time, pinned_billing_time from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost -from litellm.litellm_core_utils.llm_cost_calc.utils import get_billed_token_rates from litellm.proxy._types import ( CommonProxyErrors, CostEstimateRequest, @@ -85,9 +84,9 @@ def _extract_custom_pricing( ) -def _lookup_model_info(model: str) -> ModelInfo | None: +def _lookup_model_info(model: str, custom_llm_provider: str | None = None) -> ModelInfo | None: try: - return litellm.get_model_info(model=model) + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: return None @@ -122,7 +121,7 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: if resolved_model: verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model) custom_cost_per_token: Final = _extract_custom_pricing( - litellm_params, model_info, _lookup_model_info(str(resolved_model)) + litellm_params, model_info, _lookup_model_info(str(resolved_model), provider) ) return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token) except Exception as e: @@ -632,10 +631,9 @@ async def estimate_cost( function_id="cost-estimate", ) - # The totals, the per-token-type lines and the reported rates each resolve pricing on their - # own path. Pinning one moment keeps an off-peak window that opens mid-quote from splitting them. - billed_at: Final = current_billing_time() - with pinned_billing_time(billed_at): + # Pinning one moment keeps an off-peak window that opens mid-quote from pricing the totals on + # one side of it and the reported rates on the other. + with pinned_billing_time(current_billing_time()): # Use completion_cost which handles all the logic including margins/discounts try: cost_per_request: Final = completion_cost( @@ -653,19 +651,16 @@ async def estimate_cost( }, ) - rates: Final = get_billed_token_rates( - model=resolved_model, - custom_llm_provider=resolved_provider, - usage=usage, - custom_cost_per_token=resolved.custom_cost_per_token, - current_time=billed_at, - ) - + # The rates come back from the pricing call itself rather than a second lookup, so they are the + # ones the cost lines above billed at even when completion_cost infers a provider this endpoint + # never resolved (an unrouted "xai/grok-4" prices on xai's inclusive tier thresholds; a lookup + # here without that provider would report the sub-200k rate for a line billed above it). + rates: Final = litellm_logging_obj.billed_token_rates per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown) daily: Final = per_request.times(request.num_requests_per_day) monthly: Final = per_request.times(request.num_requests_per_month) - model_info: Final = _lookup_model_info(resolved_model) + model_info: Final = _lookup_model_info(resolved_model, resolved_provider) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 90178428018..5a19ed0277a 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -33,7 +33,6 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( CostCalculatorUtils, PromptTokensDetailsResult, TokenRates, - TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, _is_off_peak, @@ -4023,6 +4022,49 @@ def test_a_pinned_billing_time_prices_the_totals_and_the_reported_rates_at_one_m assert peak_completion_cost == pytest.approx(500 * peak_rates.output_cost_per_token) +def test_the_token_type_breakdown_carries_the_rates_it_billed_at(monkeypatch): + """Callers that report both the lines and the rates read the rates off the breakdown rather than + resolving them a second time, so the breakdown has to hand back exactly what it billed at.""" + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + breakdown = get_token_type_cost_breakdown(model="xai/tiered-model", custom_llm_provider="xai", usage=usage) + + assert breakdown.rates == get_billed_token_rates( + model="xai/tiered-model", custom_llm_provider="xai", usage=usage + ) + assert breakdown.rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert breakdown.cache_read_cost == pytest.approx(100_000 * breakdown.rates.cache_read_input_token_cost) + + +def test_the_token_type_breakdown_reports_no_rates_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + breakdown = get_token_type_cost_breakdown( + model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage + ) + + assert breakdown.rates is None + + def test_billed_token_rates_are_none_for_an_unpriced_model(): usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) @@ -4036,9 +4078,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost model="gpt-4o", custom_llm_provider="openai", usage=usage ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) @pytest.mark.parametrize( @@ -4110,9 +4150,7 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=5), ), ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map): diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index e9485f3a044..7ece35ceedf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -1139,6 +1139,45 @@ class TestEstimateCostCacheAndReasoningTokens: assert response.output_cost_per_request == pytest.approx(OUTPUT_TOKENS * response.output_cost_per_token) + @pytest.mark.asyncio + async def test_an_unrouted_model_reports_the_rates_of_the_provider_the_calculator_inferred(self, monkeypatch): + """The cost calculator infers a provider this endpoint never resolved, and the provider decides + whether a tier threshold is inclusive. xai bills a request sitting exactly on the 200k threshold + at the tier rate, so the reported rates have to be the tier's rather than the sub-tier base.""" + an_xai_model = "xai/tiered-model" + monkeypatch.setitem( + litellm.model_cost, + an_xai_model, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=an_xai_model, + input_tokens=200_000, + cache_read_input_tokens=100_000, + output_tokens=1_000, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_read_cost_per_request == pytest.approx(100_000 * response.cache_read_input_token_cost) + assert response.input_cost_per_request == pytest.approx( + 100_000 * response.input_cost_per_token + response.cache_read_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + + class TestCostEstimateRequestTokenSubsets: def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): with pytest.raises(ValidationError, match="cannot exceed input_tokens"): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 910587d3dc5..8f8a7640c08 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3736,6 +3736,62 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100 * 3e-08) +def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): + """A caller reporting the cost lines beside their per-token rates reads both off this one call. + completion_cost infers the provider, and xai's inclusive tier thresholds put a request sitting + exactly on 200k at the tier rate, which a lookup made without that inferred provider would miss. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + logging_obj = Logging( + model="xai/tiered-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="billed-rates", + function_id="f", + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + litellm.completion_cost( + completion_response=ModelResponse(model="xai/tiered-model", usage=usage), + model="xai/tiered-model", + custom_llm_provider=None, + litellm_logging_obj=logging_obj, + ) + + rates = logging_obj.billed_token_rates + assert rates is not None + assert rates.input_cost_per_token == pytest.approx(6e-6) + assert rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( + 100_000 * rates.cache_read_input_token_cost + ) + assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) + + def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing(): """ A custom-priced deployment bills cache tokens at its custom cache rates, but the From 2ece8735384e063b88ca91cf9d2f8f6d2e94b30a Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 17:05:38 +0000 Subject: [PATCH 103/157] test(e2e): lite CLI never fetches the model cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 2 + tests/e2e/other/test_cli_cost_map_e2e.py | 109 +++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 tests/e2e/other/test_cli_cost_map_e2e.py diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 814ebae2e0b..8ac4fc0762c 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -2,6 +2,8 @@ # PROMOTION NOTE: the auth cluster (~14 cells) is a candidate to promote to its own module once stable. - {id: other.auth.master_key.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "user_api_key_auth.py:1569-1588", rationale: "Master key authenticates; timing-safe compare"} - {id: other.auth.master_key.invalid_denied, module: other, tier: P0, area: auth, assertions: [invalid_denied], source: "user_api_key_auth.py:1580", rationale: "Invalid master key rejected"} +- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} +- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} - {id: other.auth.llm_chat.missing_header_denied, module: other, tier: P0, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Chat with no Authorization header is 401/403"} - {id: other.auth.llm_chat.invalid_bearer_denied, module: other, tier: P0, area: auth, assertions: [invalid_bearer_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Bearer invalid_token on chat is 401/403"} - {id: other.auth.llm_chat.no_bearer_prefix_denied, module: other, tier: P0, area: auth, assertions: [no_bearer_prefix_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Token without Bearer scheme on chat is 401/403"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py new file mode 100644 index 00000000000..6022d5a0a4c --- /dev/null +++ b/tests/e2e/other/test_cli_cost_map_e2e.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import os +import shutil +import subprocess +import threading +from collections.abc import Mapping +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final + +import pytest +from e2e_config import MASTER_KEY, PROXY_BASE_URL +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + + +def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, threading.Thread]: + class CostMapHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + with request_log.open("a", encoding="utf-8") as log_file: + log_file.write(f"{self.path}\n") + body: Final = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), CostMapHandler) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server, thread + + +def _run_lite( + args: tuple[str, ...], + server: ThreadingHTTPServer, + env: Mapping[str, str], +) -> subprocess.CompletedProcess[str]: + lite_path: Final = shutil.which("lite") + assert lite_path is not None, "the installed lite executable is required for e2e coverage" + try: + return subprocess.run( + [lite_path, *args], + capture_output=True, + text=True, + timeout=60, + env=env, + ) + finally: + server.shutdown() + server.server_close() + + +def _request_count(request_log: Path) -> int: + return request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 + + +class TestLiteCliCostMapFetch: + @pytest.mark.covers("other.cli.model_cost_map.version_skips_fetch") + def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: + request_log: Final = tmp_path / "requests.log" + server, thread = _start_cost_map_server(request_log) + source_root: Final = str(Path(__file__).resolve().parents[3]) + pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) + env: Final = { + **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "PYTHONPATH": pythonpath, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + } + try: + result: Final = _run_lite(("--version",), server, env) + finally: + thread.join(timeout=10) + + assert result.returncode == 0 + assert "LiteLLM Proxy CLI Version" in result.stdout + assert _request_count(request_log) == 0 + + @pytest.mark.covers("other.cli.model_cost_map.models_list_skips_fetch") + def test_lite_models_list_uses_proxy_not_cost_map(self, tmp_path: Path, proxy: ProxyClient) -> None: + model_names: Final = tuple(entry.model_name for entry in proxy.model_info()) + assert model_names + request_log: Final = tmp_path / "requests.log" + server, thread = _start_cost_map_server(request_log) + source_root: Final = str(Path(__file__).resolve().parents[3]) + pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) + env: Final = { + **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "PYTHONPATH": pythonpath, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + "LITELLM_PROXY_API_KEY": MASTER_KEY, + } + try: + result: Final = _run_lite(("models", "list"), server, env) + finally: + thread.join(timeout=10) + + assert result.returncode == 0 + assert result.stdout.strip() + assert any(model_name in result.stdout for model_name in model_names) + assert _request_count(request_log) == 0 From 6e71b90a888b91fe3aa281fa14da58f258b9f056 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 17:07:34 +0000 Subject: [PATCH 104/157] test(e2e): dedupe lite env setup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 4 ++-- tests/e2e/other/test_cli_cost_map_e2e.py | 29 ++++++++++-------------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 8ac4fc0762c..a7ba0dbb8cb 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -2,8 +2,6 @@ # PROMOTION NOTE: the auth cluster (~14 cells) is a candidate to promote to its own module once stable. - {id: other.auth.master_key.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "user_api_key_auth.py:1569-1588", rationale: "Master key authenticates; timing-safe compare"} - {id: other.auth.master_key.invalid_denied, module: other, tier: P0, area: auth, assertions: [invalid_denied], source: "user_api_key_auth.py:1580", rationale: "Invalid master key rejected"} -- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} -- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} - {id: other.auth.llm_chat.missing_header_denied, module: other, tier: P0, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Chat with no Authorization header is 401/403"} - {id: other.auth.llm_chat.invalid_bearer_denied, module: other, tier: P0, area: auth, assertions: [invalid_bearer_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Bearer invalid_token on chat is 401/403"} - {id: other.auth.llm_chat.no_bearer_prefix_denied, module: other, tier: P0, area: auth, assertions: [no_bearer_prefix_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "Token without Bearer scheme on chat is 401/403"} @@ -50,3 +48,5 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} +- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} +- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py index 6022d5a0a4c..62cd3d6c3f5 100644 --- a/tests/e2e/other/test_cli_cost_map_e2e.py +++ b/tests/e2e/other/test_cli_cost_map_e2e.py @@ -37,6 +37,16 @@ def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, thre return server, thread +def _lite_env(server: ThreadingHTTPServer, api_key: str | None) -> dict[str, str]: + base_env: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + return { + **base_env, + "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", + "LITELLM_PROXY_URL": PROXY_BASE_URL, + **({"LITELLM_PROXY_API_KEY": api_key} if api_key is not None else {}), + } + + def _run_lite( args: tuple[str, ...], server: ThreadingHTTPServer, @@ -66,14 +76,7 @@ class TestLiteCliCostMapFetch: def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: request_log: Final = tmp_path / "requests.log" server, thread = _start_cost_map_server(request_log) - source_root: Final = str(Path(__file__).resolve().parents[3]) - pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) - env: Final = { - **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, - "PYTHONPATH": pythonpath, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - } + env: Final = _lite_env(server, None) try: result: Final = _run_lite(("--version",), server, env) finally: @@ -89,15 +92,7 @@ class TestLiteCliCostMapFetch: assert model_names request_log: Final = tmp_path / "requests.log" server, thread = _start_cost_map_server(request_log) - source_root: Final = str(Path(__file__).resolve().parents[3]) - pythonpath: Final = os.pathsep.join(filter(None, (source_root, os.environ.get("PYTHONPATH")))) - env: Final = { - **{key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"}, - "PYTHONPATH": pythonpath, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - "LITELLM_PROXY_API_KEY": MASTER_KEY, - } + env: Final = _lite_env(server, MASTER_KEY) try: result: Final = _run_lite(("models", "list"), server, env) finally: From 996ee5635ab79854992661279c2e8db793b3dda4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:35:50 -0700 Subject: [PATCH 105/157] perf(proxy): pipeline spend counter increments into one Redis call per request (#40371) * perf(proxy): pipeline spend counter increments into one redis call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply surviving spend increments before raising scope error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format spend counter helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): settle inner spend counter gathers and fall back per key on pipeline failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): suppress BLE001 on pipeline fallback catch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): invalidate all batched spend counters on pipeline failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 331 +++++++---- .../proxy/proxy_server/test_spend_counters.py | 538 ++++++++++-------- .../test_budget_reservation_redis_failure.py | 9 + .../proxy/test_budget_reservation.py | 9 + tests/test_litellm/proxy/test_proxy_server.py | 94 ++- 5 files changed, 609 insertions(+), 372 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7aed0553b4b..c241e66049b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17,6 +17,7 @@ import time import traceback import warnings from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -131,6 +132,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -2703,6 +2705,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False +@dataclass(frozen=True, slots=True) +class _PendingSpendIncrement: + counter_key: str + increment: float + + async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2737,7 +2745,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> None: + async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2748,30 +2756,29 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - if key_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=key_counter_key, - source_cache_key=hashed_token, - increment=cost, + key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if key_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=key_counter_key, + source_cache_key=hashed_token, + increment=cost, + ), ) - - key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) - if key_obj is None: - return - key_budget_limits = getattr(key_obj, "budget_limits", None) or ( - key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(key_budget_limits, str): - key_budget_limits = json.loads(key_budget_limits) - if not isinstance(key_budget_limits, list): - return - for window in key_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + + async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + key_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" key_window_start = get_budget_window_start(window) - if key_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, @@ -2779,6 +2786,9 @@ async def increment_spend_counters( window_start=key_window_start, increment=cost, ) + if key_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, @@ -2788,33 +2798,48 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_scope(scope_team_id: str) -> None: - team_counter_key: Final = f"spend:team:{scope_team_id}" - if team_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=team_counter_key, - source_cache_key=f"team_id:{scope_team_id}", - increment=cost, - ) - - team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") - if team_obj is None: - return - team_budget_limits = getattr(team_obj, "budget_limits", None) or ( - team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is None: + return key_pending + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(team_budget_limits, str): - team_budget_limits = json.loads(team_budget_limits) - if not isinstance(team_budget_limits, list): - return - for window in team_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if not isinstance(key_budget_limits, list): + return key_pending + window_pending: Final = await asyncio.gather( + *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True + ) + return key_pending + tuple(item for item in window_pending if item is not None) + + async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + team_counter_key: Final = f"spend:team:{scope_team_id}" + team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if team_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=team_counter_key, + source_cache_key=f"team_id:{scope_team_id}", + increment=cost, + ), + ) + ) + + async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + team_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" team_window_start = get_budget_window_start(window) - if team_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, @@ -2822,6 +2847,9 @@ async def increment_spend_counters( window_start=team_window_start, increment=cost, ) + if team_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.TEAM, entity_id=scope_team_id, @@ -2831,25 +2859,47 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: + team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + if team_obj is None: + return team_pending + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if not isinstance(team_budget_limits, list): + return team_pending + window_pending: Final = await asyncio.gather( + *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True + ) + return team_pending + tuple(item for item in window_pending if item is not None) + + async def _team_member_scope( + scope_user_id: str, scope_team_id: str + ) -> tuple[_PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=team_member_counter_key, - source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", + increment=cost, + ), ) - async def _user_scope(scope_user_id: str) -> None: + async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=user_counter_key, - source_cache_key=scope_user_id, - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=user_counter_key, + source_cache_key=scope_user_id, + increment=cost, + ), ) scope_coros: Final = tuple( @@ -2859,7 +2909,7 @@ async def increment_spend_counters( _team_scope(team_id) if team_id is not None else None, _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, _user_scope(user_id) if user_id is not None else None, - _increment_end_user_and_tag_spend_counters( + _prepare_end_user_and_tag_spend_increments( end_user_id=end_user_id, tags=tags, response_cost=cost, @@ -2867,14 +2917,14 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, - _increment_model_access_group_spend_counters( + _prepare_model_access_group_spend_increments( model_access_groups=model_access_groups, response_cost=cost, reserved_counter_keys=reserved_counter_keys, ) if model_access_groups else None, - _increment_org_spend_counter( + _prepare_org_spend_increment( org_id=org_id, response_cost=cost, reserved_counter_keys=reserved_counter_keys, @@ -2889,7 +2939,20 @@ async def increment_spend_counters( # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)] + scope_errors: Final = tuple( + item + for scope in scope_results + for item in (scope if isinstance(scope, tuple) else (scope,)) + if isinstance(item, BaseException) + ) + pending: Final = tuple( + item + for scope in scope_results + if not isinstance(scope, BaseException) + for item in scope + if not isinstance(item, BaseException) + ) + await _apply_spend_counter_increments(pending=pending) if scope_errors: raise scope_errors[0] @@ -2932,41 +2995,49 @@ async def _reconcile_budget_reservation_for_counter_update( return reserved_counter_keys -async def _increment_end_user_and_tag_spend_counters( +async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: - if end_user_id is not None: - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=end_user_cache_key(end_user_id), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) - - if tags is None: - return - - seen_tags: Final[set[str]] = set() - for tag_name in tags: - if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: - continue - seen_tags.add(tag_name) - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:tag:{tag_name}", - source_cache_key=tag_cache_key(tag_name), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) +) -> tuple[_PendingSpendIncrement | BaseException, ...]: + unique_tags: Final = ( + tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () + ) + results: Final = await asyncio.gather( + *( + coro + for coro in ( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + if end_user_id is not None + else None, + *( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for tag_name in unique_tags + ), + ) + if coro is not None + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_model_access_group_spend_counters( +async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -2980,55 +3051,63 @@ async def _increment_model_access_group_spend_counters( unique_groups: Final = tuple( dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) ) - for group in unique_groups: - await _init_and_increment_unreserved_spend_counter( - counter_key=model_access_group_spend_counter_key(group), - source_cache_key=model_access_group_cache_key(group), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) + results: Final = await asyncio.gather( + *( + _prepare_unreserved_spend_counter_increment( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for group in unique_groups + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_org_spend_counter( +async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement, ...]: if org_id is None: - return + return () - await _init_and_increment_unreserved_spend_counter( + pending: Final = await _prepare_unreserved_spend_counter_increment( counter_key=f"spend:org:{org_id}", source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"], increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) + return (pending,) if pending is not None else () -async def _init_and_increment_unreserved_spend_counter( +async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> None: +) -> _PendingSpendIncrement | None: if counter_key in reserved_counter_keys: - return + return None - await _init_and_increment_spend_counter( + return await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, increment=increment, ) -async def _init_and_increment_spend_counter( +async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -): +) -> _PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet - set, then atomically increment in both in-memory and Redis. + set, then return the pending increment for the caller to apply in one + pipelined Redis call. On first access per pod: 1. Check spend_counter_cache (in-memory -> Redis via DualCache) @@ -3040,13 +3119,13 @@ async def _init_and_increment_spend_counter( the counter as absent and seed it. Using increment means the worst case is over-counting (conservative, blocks slightly early) rather than under-counting (would allow overspend). - 4. Increment atomically (both in-memory + Redis) + 4. Increment is returned for the caller to apply via pipeline """ await _ensure_spend_counter_initialized( counter_key=counter_key, source_cache_key=source_cache_key, ) - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3098,20 +3177,20 @@ async def _enqueue_window_spend_row_update( ) -async def _init_and_increment_window_spend_counter( +async def _prepare_window_spend_counter_increment( counter_key: str, entity_type: str, entity_id: str, window_duration: str | None, window_start: datetime | None, increment: float, -): +) -> _PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", counter_key, ) - return + return None initialized: Final = await _ensure_window_spend_counter_initialized( counter_key=counter_key, @@ -3121,8 +3200,8 @@ async def _init_and_increment_window_spend_counter( window_start=window_start, ) if initialized is False: - return - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return None + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3255,6 +3334,32 @@ async def _invalidate_spend_counter(counter_key: str): ) +async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: + if not pending: + return + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + for item in pending: + await spend_counter_cache.async_increment_cache( + key=item.counter_key, + value=item.increment, + refresh_ttl=True, + ) + return + ttl: Final = redis_cache.get_ttl() + increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation] + RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) + for item in pending + ] + try: + results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) + except Exception: + await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) + raise + for item, current_value in zip(pending, results or ()): + spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + + async def update_cache( token: str | None, user_id: str | None, diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 86e97a334df..c343652efd9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -4,11 +4,12 @@ Pins covered: - ``get_current_spend`` - ``increment_spend_counters`` - ``_reconcile_budget_reservation_for_counter_update`` -- ``_increment_end_user_and_tag_spend_counters`` -- ``_increment_org_spend_counter`` -- ``_init_and_increment_unreserved_spend_counter`` -- ``_init_and_increment_spend_counter`` -- ``_init_and_increment_window_spend_counter`` +- ``_prepare_end_user_and_tag_spend_increments`` +- ``_prepare_org_spend_increment`` +- ``_prepare_unreserved_spend_counter_increment`` +- ``_prepare_spend_counter_increment`` +- ``_prepare_window_spend_counter_increment`` +- ``_apply_spend_counter_increments`` - ``_ensure_spend_counter_initialized`` - ``_get_source_cache_base_spend`` - ``_ensure_window_spend_counter_initialized`` @@ -48,9 +49,7 @@ def _make_spend_counter_cache( cache.in_memory_cache.delete_cache = MagicMock() if with_redis: cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock( - return_value=redis_get_value, side_effect=redis_get_side_effect - ) + cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) cache.redis_cache.async_increment = AsyncMock( return_value=redis_increment_value, side_effect=redis_increment_side_effect, @@ -58,6 +57,8 @@ def _make_spend_counter_cache( cache.redis_cache.async_delete_cache = AsyncMock() cache.redis_cache.async_set_cache = AsyncMock() cache.redis_cache.async_set_max = AsyncMock() + cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) + cache.redis_cache.get_ttl = MagicMock(return_value=None) else: cache.redis_cache = None cache.async_increment_cache = AsyncMock(return_value=redis_increment_value) @@ -70,9 +71,7 @@ def _make_spend_counter_cache( def _make_user_api_key_cache(get_value=None, get_side_effect=None): cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value=get_value, side_effect=get_side_effect - ) + cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect) cache.async_set_cache_pipeline = AsyncMock() return cache @@ -109,9 +108,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch ) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=99.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -136,9 +133,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) @pytest.mark.asyncio @@ -169,9 +164,7 @@ async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0) assert result == 2.0 assert from_db.await_count == 0 @@ -210,12 +203,8 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - first = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) - second = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) + first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) + second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) assert first == 12.0 assert second == 12.0 @@ -336,9 +325,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -379,9 +366,7 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) @pytest.mark.asyncio @@ -393,9 +378,7 @@ async def test_get_current_spend_floors_window_against_logs_when_row_stale(monke window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) fake_prisma = _make_window_spend_prisma( - row=SimpleNamespace( - window_start=window_start - timedelta(days=7), spend=999.0 - ), + row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0), spend_logs_total=15.0, ) fake_cache = _make_spend_counter_cache(redis_get_value=2.0) @@ -423,21 +406,13 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) with pytest.raises(HTTPException) as exc: - await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert exc.value.status_code == 503 @@ -445,18 +420,12 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "general_settings", {}) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -466,13 +435,9 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa authoritative, so an under-budget request is admitted normally.""" fake_cache = _make_spend_counter_cache(redis_get_value=1.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -481,16 +446,10 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) result = await ps.get_current_spend( counter_key="spend:end_user:e1", @@ -508,9 +467,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa re-checks the authoritative DB and enforces against it.""" fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -532,9 +489,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -543,9 +498,7 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): async def _fake_coalesced(**kwargs): return None - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) await ps.increment_spend_counters( token="hashed-tok", @@ -554,25 +507,36 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): response_cost=5.0, ) + pipeline = fake_cache.redis_cache.async_increment_pipeline + pipeline.assert_awaited_once() + increment_list = pipeline.await_args.kwargs["increment_list"] + assert {op["key"] for op in increment_list} == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + } + assert all(op["increment_value"] == 5.0 for op in increment_list) observed = { "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, + "pipeline_calls": pipeline.await_count, "user_cache_used": fake_user_cache.async_get_cache.called, } assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 4, + "redis_increment_called": False, + "pipeline_calls": 1, "user_cache_used": True, } class _ConcurrencyProbe: - """Stand-in for redis_cache.async_increment that pins concurrency. + """Stand-in for redis_cache.async_get_cache that pins concurrency. - Each call registers itself as in-flight and blocks on ``release`` until the - test lets it proceed. ``all_arrived`` fires once ``expected`` distinct scope - increments are simultaneously suspended here, which can only happen if the - per-scope increments are gathered rather than awaited one after another. + Each warm-check read registers itself as in-flight and blocks on ``release`` + until the test lets it proceed. ``all_arrived`` fires once ``expected`` + distinct scope warm-checks are simultaneously suspended here, which can only + happen if the per-scope prepares are gathered rather than awaited one after + another. """ def __init__(self, expected_concurrency: int): @@ -581,36 +545,45 @@ class _ConcurrencyProbe: self.max_in_flight = 0 self.all_arrived = asyncio.Event() self.release = asyncio.Event() - self.values: dict[str, float] = {} + self.keys: list[str] = [] - async def async_increment(self, *, key, value, refresh_ttl=True): + async def async_get_cache(self, *, key, **kwargs): self.in_flight += 1 self.max_in_flight = max(self.max_in_flight, self.in_flight) + self.keys.append(key) if self.in_flight >= self.expected: self.all_arrived.set() if not self.release.is_set(): await self.release.wait() self.in_flight -= 1 - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] + return 1.0 @pytest.mark.asyncio async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): """The six independent scopes (key, team, team_member, user, end_user+tags, - org) must be incremented concurrently. The probe only fires once all six are - suspended in async_increment at the same time, which is impossible if the + org) must prepare their increments concurrently. The probe only fires once + all eight warm-check reads (one per counter: 6 scopes + 2 tags) are + suspended in async_get_cache at the same time, which is impossible if the awaits are chained sequentially.""" - probe = _ConcurrencyProbe(expected_concurrency=6) - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = probe.async_increment + probe = _ConcurrencyProbe(expected_concurrency=8) + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = probe.async_get_cache + recorded: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results + + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) task = asyncio.create_task( ps.increment_spend_counters( @@ -630,16 +603,16 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): probe.release.set() await task pytest.fail( - "scope increments did not run concurrently; sequential awaits " - f"detected (peak in-flight was {probe.max_in_flight}, expected 6)" + "scope prepares did not run concurrently; sequential awaits " + f"detected (peak in-flight was {probe.max_in_flight}, expected 8)" ) - assert probe.in_flight == 6 - assert probe.max_in_flight == 6 + assert probe.in_flight == 8 + assert probe.max_in_flight == 8 probe.release.set() await task - assert probe.values == { + assert recorded == { "spend:key:hashed-tok": 5.0, "spend:team:t1": 5.0, "spend:team_member:u1:t1": 5.0, @@ -659,26 +632,25 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) import litellm.proxy.spend_tracking.budget_reservation as br reserved = {"spend:key:hashed-tok", "spend:org:org1"} - monkeypatch.setattr( - br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)) - ) + monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock()) recorded: dict[str, float] = {} - async def _record_increment(*, key, value, refresh_ttl=True): - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _record_increment + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} await ps.increment_spend_counters( @@ -708,27 +680,46 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ): """A failure in one scope must propagate to the caller (so it can invalidate reserved counters) while every other scope still settles rather than being - left as an orphaned background task, and the reservation is not finalized.""" - recorded: dict[str, float] = {} + left as an orphaned background task, and the reservation is not finalized. + The surviving scopes' increments are still applied in the single pipeline: + dropping them would under-count spend, the unsafe direction for budget + enforcement.""" + warmed_keys: list[str] = [] - async def _increment(*, key, value, refresh_ttl=True): + async def _warm_check(*, key, **kwargs): + warmed_keys.append(key) if key == "spend:team:t1": - raise RuntimeError("redis increment failed") - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + raise RuntimeError("redis get failed") + return 1.0 - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _increment + async def _reseed_fails(*, counter_key, **kwargs): + if counter_key == "spend:team:t1": + raise RuntimeError("reseed failed") + + applied: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + applied[op["key"]] = op["increment_value"] + results.append(op["increment_value"]) + return results + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check) + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) + ps.SpendCounterReseed, + "coalesced", + AsyncMock(side_effect=_reseed_fails), ) reservation = {"finalized": False} - with pytest.raises(RuntimeError, match="redis increment failed"): + with pytest.raises(RuntimeError, match="reseed failed"): await ps.increment_spend_counters( token="hashed-tok", team_id="t1", @@ -741,7 +732,19 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ) assert reservation["finalized"] is False - assert recorded == { + # every sibling scope settled (its warm-check ran) before the error propagated + assert set(warmed_keys) == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:org:org1", + } + # the surviving scopes' increments were still applied, in one pipeline call + fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once() + assert applied == { "spend:key:hashed-tok": 5.0, "spend:team_member:u1:t1": 5.0, "spend:user:u1": 5.0, @@ -749,6 +752,7 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ "spend:tag:a": 5.0, "spend:org:org1": 5.0, } + fake_cache.redis_cache.async_increment.assert_not_awaited() @pytest.mark.asyncio @@ -772,6 +776,108 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( assert reservation == {"finalized": True} assert fake_cache.redis_cache.async_increment.called is False + fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( + monkeypatch, +): + """Every scope's increment must go out in a single async_increment_pipeline + call, not one INCRBYFLOAT round-trip per scope.""" + counter_cache = ps.DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + + async def _pipeline(increment_list, **_): + return [1.5] * len(increment_list) + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline) + fake_redis.async_increment = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + fake_redis.async_increment_pipeline.assert_awaited_once() + assert fake_redis.async_increment.await_count == 0 + increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"] + expected_keys = { + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + } + assert {op["key"] for op in increment_list} == expected_keys + assert all(op["increment_value"] == 0.5 for op in increment_list) + for key in expected_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import MaxConnectionsError + + counter_cache = ps.DualCache() + pending_keys = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + ) + for key in pending_keys: + counter_cache.in_memory_cache.set_cache(key=key, value=1.0) + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment = AsyncMock() + fake_redis.async_delete_cache = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + with pytest.raises(MaxConnectionsError): + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + assert fake_redis.async_increment.await_count == 0 + deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} + assert deleted_keys == set(pending_keys) + for key in pending_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) is None # --------------------------------------------------------------------------- @@ -781,9 +887,7 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( @pytest.mark.asyncio async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none(): - result = await ps._reconcile_budget_reservation_for_counter_update( - budget_reservation=None, response_cost=1.0 - ) + result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0) assert result == set() @@ -818,179 +922,151 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat # --------------------------------------------------------------------------- -# _increment_end_user_and_tag_spend_counters +# _prepare_end_user_and_tag_spend_increments # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag( +async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=3.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id="eu1", tags=["a", "b", "a", "", None], response_cost=3.0, reserved_counter_keys=set(), ) - observed = { - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - "called": fake_cache.redis_cache.async_increment.called, - } - assert normalize(observed) == { - "increment_calls": 3, - "in_memory_set_calls": 3, - "called": True, + assert {item.counter_key for item in pending} == { + "spend:end_user:eu1", + "spend:tag:a", + "spend:tag:b", } + assert all(item.increment == 3.0 for item in pending) @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_invalid_input_noop( +async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id=None, tags=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _increment_org_spend_counter +# _prepare_org_spend_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=10.0 - ) +async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id="org-1", response_cost=10.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[ - "key" - ], - } - assert normalize(observed) == { - "increment_called": True, - "increment_calls": 1, - "counter_key_arg": "spend:org:org-1", - } + assert len(pending) == 1 + assert pending[0].counter_key == "spend:org:org-1" + assert pending[0].increment == 10.0 @pytest.mark.asyncio -async def test_increment_org_spend_counter_no_org_is_noop_invalid_id(monkeypatch): +async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _init_and_increment_unreserved_spend_counter +# _prepare_unreserved_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys( +async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:x", source_cache_key="tag:x", increment=1.0, reserved_counter_keys={"spend:tag:x"}, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved( +async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=2.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None) fake_user_cache = _make_user_api_key_cache() + reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:y", source_cache_key="tag:y", increment=2.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - "reseed_consulted": True, - } - assert observed == { - "increment_called": True, - "redis_get_called": True, - "reseed_consulted": True, - } + assert pending is not None + assert pending.counter_key == "spend:tag:y" + assert pending.increment == 2.0 + assert fake_cache.redis_cache.async_get_cache.called is True + assert reseed.called is True # --------------------------------------------------------------------------- -# _init_and_increment_spend_counter +# _prepare_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=11.0, redis_increment_value=14.0 - ) +async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=11.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -998,12 +1074,14 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_spend_counter( + pending = await ps._prepare_spend_counter_increment( counter_key="spend:key:k", source_cache_key="k", increment=3.0, ) + assert pending.counter_key == "spend:key:k" + assert pending.increment == 3.0 observed = { "reseed_called": reseed.called, "increment_called": fake_cache.redis_cache.async_increment.called, @@ -1011,23 +1089,21 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa } assert normalize(observed) == { "reseed_called": False, - "increment_called": True, + "increment_called": False, "in_memory_seeded_from_redis": True, } # --------------------------------------------------------------------------- -# _init_and_increment_window_spend_counter +# _prepare_window_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_increments_when_initialized( +async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=0.0, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( @@ -1036,7 +1112,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ AsyncMock(return_value=0.0), ) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1045,26 +1121,19 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ increment=5.0, ) - observed = { - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - } - assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 1, - "in_memory_set_calls": 2, - } + assert pending is not None + assert pending.counter_key == "spend:key:k:window:1d" + assert pending.increment == 5.0 @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_missing_window_start_invalid_skips( +async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1073,6 +1142,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva increment=5.0, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @@ -1114,16 +1184,12 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=7.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._ensure_spend_counter_initialized( counter_key="spend:user:u", @@ -1163,9 +1229,7 @@ async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch fake_user_cache.async_get_cache = AsyncMock(side_effect=_get) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - result = await ps._get_source_cache_base_spend( - source_cache_key=["miss", "hit-obj", "miss2"] - ) + result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"]) observed = { "result": result, @@ -1294,9 +1358,7 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=4.0 - ) + result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) observed = { "result": result, @@ -1314,15 +1376,11 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_increment_side_effect=RuntimeError("incr fail") - ) + fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) with pytest.raises(RuntimeError): - await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=1.0 - ) + await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) assert fake_cache.in_memory_cache.delete_cache.called is True assert fake_cache.redis_cache.async_delete_cache.called is True @@ -1343,9 +1401,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) observed = { "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[ - "key" - ], + "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], } assert normalize(observed) == { "in_memory_delete_called": True, @@ -1357,9 +1413,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("redis down") - ) + fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._invalidate_spend_counter(counter_key="spend:key:k") diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py index c123eeeed36..6165af4920d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -41,6 +41,15 @@ class _FlakyRedisCache: self._store[key] = float(value) return True + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs): + return None + @pytest.mark.asyncio async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..40ebc03781c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2220,6 +2220,15 @@ class _ExpiringRedisCache: async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs) -> None: + return None + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54579e6cb7c..b0f38978727 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7749,7 +7749,7 @@ async def test_increment_spend_counters_team_and_member(): @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(): +async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(): """When the Redis counter is missing, the reseed path reads the authoritative spend from the DB (not a stale cache), so the next increment continues from the correct base value.""" @@ -7762,8 +7762,17 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( recorded_increments.append({"key": key, "value": value, "ttl": ttl}) return value + async def record_pipeline(increment_list, **kwargs): + results = [] + for op in increment_list: + await record_increment(key=op["key"], value=op["increment_value"], ttl=op["ttl"]) + results.append(op["increment_value"]) + return results + fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=record_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis @@ -7782,7 +7791,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( stale_cache.in_memory_cache.set_cache(key="team_id:team-9", value=stale_team) import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) orig_user, orig_counter, orig_prisma = ( ps.user_api_key_cache, @@ -7793,11 +7805,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key="spend:team:team-9", source_cache_key="team_id:team-9", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. @@ -7976,7 +7989,10 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): @pytest.mark.asyncio async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) @@ -7992,7 +8008,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", @@ -8000,6 +8016,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8015,7 +8032,10 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): @pytest.mark.asyncio async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:team:team-stale-local" @@ -8037,6 +8057,15 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -8055,11 +8084,12 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.prisma_client = fake_prisma ps.user_api_key_cache = DualCache() try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="team_id:team-stale-local", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. @@ -8074,7 +8104,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-stale-local:window:1h" @@ -8097,6 +8130,15 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8111,7 +8153,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", @@ -8119,6 +8161,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8138,7 +8181,10 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-concurrent-seed:window:1h" @@ -8161,6 +8207,15 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) fake_redis.async_set_cache = AsyncMock(return_value=False) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8175,7 +8230,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", @@ -8183,6 +8238,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_redis.async_set_cache.assert_awaited_once_with( key=counter_key, @@ -8199,7 +8255,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() @pytest.mark.asyncio async def test_window_spend_counter_skips_invalid_window_start(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import _prepare_window_spend_counter_increment counter_cache = DualCache() @@ -8208,7 +8264,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): orig_counter = ps.spend_counter_cache ps.spend_counter_cache = counter_cache try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", @@ -8216,6 +8272,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): window_start=None, increment=0.5, ) + assert pending is None assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: @@ -8279,6 +8336,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) + return ps._PendingSpendIncrement( + counter_key=kwargs["counter_key"], increment=kwargs["increment"] + ) import litellm.proxy.proxy_server as ps @@ -8287,7 +8347,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): ps.user_api_key_cache = DualCache() try: with patch( - "litellm.proxy.proxy_server._init_and_increment_spend_counter", + "litellm.proxy.proxy_server._prepare_spend_counter_increment", new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), ): await increment_spend_counters( @@ -8620,7 +8680,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): async def test_concurrent_read_and_write_paths_share_one_db_query(): """ The read path (`get_current_spend`) and the write path - (`_init_and_increment_spend_counter`) both reseed cold counters from + (`_prepare_spend_counter_increment`) both reseed cold counters from the DB. They must share the per-counter lock so a concurrent pre-call enforcement read and post-call increment for the same counter collapse to one DB query, not two. @@ -8629,7 +8689,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ( - _init_and_increment_spend_counter, + _prepare_spend_counter_increment, get_current_spend, ) @@ -8683,7 +8743,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): try: results = await _asyncio.gather( get_current_spend(counter_key=counter_key, fallback_spend=0.0), - _init_and_increment_spend_counter( + _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="ignored", increment=1.5, From 7c6e33ef70889445fab5650b1490ddf3c0448a11 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:38:26 -0700 Subject: [PATCH 106/157] fix(proxy): ignore team_id="" on /key/update so team-less keys can be updated and imported (#40421) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 8 ++++++ .../test_key_management_endpoints.py | 26 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index df7ee8f508c..0e2fd50ac9f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3,6 +3,7 @@ import json import os from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple import httpx @@ -1294,6 +1295,13 @@ class UpdateKeyRequest(KeyRequestBase): rotation_interval: str | None = None organization_id: str | None = None + @model_validator(mode="before") + @classmethod + def drop_blank_team_id(cls, values: object) -> object: + if isinstance(values, Mapping) and values.get("team_id") == "": + return MappingProxyType({k: v for k, v in values.items() if k != "team_id"}) + return values + @field_validator("organization_id", mode="before") @classmethod def treat_cleared_organization_id_as_unset(cls, v: object) -> object: 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 a873a367eab..d2aeba18f7d 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 @@ -17975,6 +17975,32 @@ def test_key_request_blank_organization_id_is_unset(): assert UpdateKeyRequest(key="sk-1", organization_id="org-1").organization_id == "org-1" +def test_update_key_request_blank_team_id_is_not_a_team_change(): + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + is_different_team, + ) + + blank = UpdateKeyRequest(key="sk-1", team_id="", key_alias="renamed") + assert blank.team_id is None + assert "team_id" not in blank.model_dump(exclude_unset=True) + assert blank.model_dump(exclude_unset=True) == {"key": "sk-1", "key_alias": "renamed"} + assert is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed")) is False + assert ( + is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed", team_id="team-1")) + is False + ) + assert "team_id" in UpdateKeyRequest(key="sk-1", team_id=None).model_dump(exclude_unset=True) + assert UpdateKeyRequest(key="sk-1", team_id="team-1").team_id == "team-1" + assert ( + is_different_team( + data=UpdateKeyRequest(key="sk-1", team_id="team-1"), + existing_key_row=LiteLLM_VerificationToken(token="hashed"), + ) + is True + ) + + def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): """key_generation_check with team_id="" must take the personal-key path instead of failing the team lookup with "Unable to find team object" (LIT-3925).""" From 699ae63b2a3270ab8b1068f8beba4a3bf71543b9 Mon Sep 17 00:00:00 2001 From: Clement Date: Thu, 10 Sep 2026 01:47:34 +0800 Subject: [PATCH 107/157] feat(router): support percentile-based TTFT routing (#40352) * feat(router): support percentile-based TTFT routing * fix(router): apply routing_strategy_args updates to the live selector Runtime routing_strategy_args updates (config reload, update_settings) only rebuilt the strategy selector when routing_strategy itself changed, so a newly added ttft_percentile sat unused until the proxy restarted. Also drops a comment that only restated the code it sat above. Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB * refactor(router): drop unreachable empty-samples guard in percentile latency _percentile_latency is only called behind use_ttft, which already requires a non-empty ttft sample list, so the early return was dead code and the one line Codecov flagged as uncovered on this patch. Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB * test(router): cover the no-selector path of a routing_strategy_args update simple-shuffle has no selector attribute to re-link, so the early return guards a setattr with a None attribute name. Dropping the guard makes the new test fail with "attribute name must be string, not 'NoneType'". Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB * fix(test): assert ValidationError on out-of-range ttft_percentile pytest.raises(ValueError) tripped PT011 for being too broad. Pydantic raises ValidationError for the gt/le constraint, so naming it satisfies the rule and pins the assertion to the constraint under test. Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB * fix(router): drop Final from a per-deployment loop variable basedpyright rejects "A Final variable cannot be assigned within a loop", which pushed reportGeneralTypeIssues one over its budget. selected_latency is rebound each iteration, so it matches its unannotated neighbours in the same loop. Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB * test(router): exempt _apply_updated_routing_strategy_args from the name scan The scan only reads test files with "router" in the filename, so it cannot see the update_settings tests in router_strategy/test_lowest_latency.py. Calling the private helper directly would test structure rather than behaviour, so it joins the existing entries ignored for the same reason. Claude-Session: https://claude.ai/code/session_01PmqjhFYcUh6vA72d8W9gdB --- litellm/router.py | 47 ++++- litellm/router_strategy/lowest_latency.py | 21 ++- .../router_code_coverage.py | 1 + .../router_strategy/test_lowest_latency.py | 174 +++++++++++++++++- 4 files changed, 225 insertions(+), 18 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 934a4ac86a9..2fbdf541487 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1317,6 +1317,43 @@ class Router: if isinstance(litellm.input_callback, list): litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + def _apply_updated_routing_strategy_args(self) -> None: + """ + Re-link the default group's selector to the current `routing_strategy_args`. + + Selectors freeze their `RoutingArgs` at construction, so a runtime args + update would otherwise keep serving the boot-time values until restart. + Latency/usage state survives the rebuild: it lives in the shared router + cache, not on the selector. + """ + strategy: Final = self._normalize_strategy(self.routing_strategy) + if strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(self, self.routing_strategy_args) + return + + attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") + current: Final = getattr(self, attr, None) if attr is not None else None + if attr is None or current is None: + return + + try: + rebuilt: Final = self._build_strategy_selector( + strategy=strategy or "", + routing_strategy_args=self.routing_strategy_args, + ) + except (TypeError, ValidationError): + verbose_router_logger.exception( + "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", + self.routing_strategy_args, + strategy, + ) + return + + self._unregister_router_selectors((current,)) + setattr(self, attr, rebuilt) + def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict): verbose_router_logger.info("Routing strategy: %s", routing_strategy) self._validate_routing_strategy(routing_strategy) @@ -11847,7 +11884,7 @@ class Router: _existing_router_settings: Final = self.get_settings() rebuild_routing_groups = False - relink_lar1_from_args = False + routing_args_updated = False for var in kwargs: if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: if var in _int_settings: @@ -11886,15 +11923,13 @@ class Router: ) rebuild_routing_groups = True elif var == "routing_strategy_args": - relink_lar1_from_args = True + routing_args_updated = True setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) - if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self, self.routing_strategy_args) + if routing_args_updated: + self._apply_updated_routing_strategy_args() if rebuild_routing_groups: self._init_routing_groups(self._routing_groups_input) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 805d4ff9080..e902192811c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -3,8 +3,11 @@ import random from collections.abc import Sequence from datetime import datetime, timedelta +from math import ceil from typing import TYPE_CHECKING, Any, Final +from pydantic import Field + import litellm from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache @@ -24,6 +27,7 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: float = 1 * 60 * 60 # 1 hour lowest_latency_buffer: float = 0 max_latency_list_size: int = 10 + ttft_percentile: float | None = Field(default=None, gt=0, le=1) def _average_latency(samples: Sequence[float]) -> float: @@ -32,6 +36,12 @@ def _average_latency(samples: Sequence[float]) -> float: return sum(samples) / len(samples) +def _percentile_latency(samples: Sequence[float], percentile: float) -> float: + values: Final = sorted(samples) + index: Final = ceil(len(values) * percentile) - 1 + return values[index] + + def _ttft_seconds(elapsed: timedelta | float) -> float: if isinstance(elapsed, timedelta): return elapsed.total_seconds() @@ -427,14 +437,17 @@ class LowestLatencyLoggingHandler(CustomLogger): item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) - # get average latency or average ttft (depending on streaming/non-streaming) use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True and len(item_ttft_latency) > 0 ) - average_latency = _average_latency(item_ttft_latency if use_ttft else item_latency) + selected_latency = ( + _percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile) + if use_ttft and self.routing_args.ttft_percentile is not None + else _average_latency(item_ttft_latency if use_ttft else item_latency) + ) # -------------- # # Debugging Logic @@ -443,7 +456,7 @@ class LowestLatencyLoggingHandler(CustomLogger): # this helps a user to debug why the router picked a specfic deployment # _deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "") if _deployment_api_base is not None: - _latency_per_deployment[_deployment_api_base] = average_latency + _latency_per_deployment[_deployment_api_base] = selected_latency # -------------- # # End of Debugging Logic # -------------- # @@ -453,7 +466,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ): # if user passed in tpm / rpm in the model_list continue else: - potential_deployments.append((_deployment, average_latency)) + potential_deployments.append((_deployment, selected_latency)) if len(potential_deployments) == 0: return None diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 0af29f069c6..a11f015743b 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -87,6 +87,7 @@ ignored_function_names = [ "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py + "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) ] diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index 812d7bbff32..1a8614e3fca 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -8,11 +8,12 @@ import json from datetime import datetime, timedelta import pytest - +from pydantic import ValidationError import litellm from litellm.caching.caching import DualCache -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.router import Router +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler, RoutingArgs DEPLOYMENT_ID = "9876" KWARGS = { @@ -58,9 +59,9 @@ def test_sync_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(2.0) # the exact failure mode from production: redis cache sync json.dumps json.dumps({"latency": latencies}) @@ -84,9 +85,9 @@ async def test_async_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(3.0) json.dumps({"latency": latencies}) @@ -292,6 +293,85 @@ async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_worke assert picked["model_info"]["id"] == FAST_TTFT_ID +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("ttft_percentile", "first_samples", "second_samples", "expected_id"), + [ + (None, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], SLOW_TTFT_ID), + (0.5, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], FAST_TTFT_ID), + (0.9, [0.1, 0.1, 0.1, 0.1, 1.5], [0.3, 0.3, 0.3, 0.3, 0.3], SLOW_TTFT_ID), + ], + ids=["default_average", "p50", "p90"], +) +async def test_streaming_ttft_ranking_percentile( + sync_mode: bool, + ttft_percentile: float | None, + first_samples: list[float], + second_samples: list[float], + expected_id: str, +): + cache = DualCache() + routing_args = {} if ttft_percentile is None else {"ttft_percentile": ttft_percentile} + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": first_samples}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": second_samples}, + }, + ) + + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == expected_id + + +@pytest.mark.parametrize("ttft_percentile", [0, -0.1, 1.1]) +def test_ttft_percentile_validation(ttft_percentile: float): + with pytest.raises(ValidationError): + RoutingArgs(ttft_percentile=ttft_percentile) + + +@pytest.mark.parametrize("ttft_percentile", [0.5, 0.9, 0.95, 1.0]) +def test_ttft_percentile_accepts_valid_values(ttft_percentile: float): + assert RoutingArgs(ttft_percentile=ttft_percentile).ttft_percentile == ttft_percentile + + +@pytest.mark.asyncio +async def test_ttft_percentile_does_not_change_non_streaming_routing(): + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"ttft_percentile": 0.9}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"latency": [1.0], "time_to_first_token_seconds": [0.1]}, + SLOW_TTFT_ID: {"latency": [0.2], "time_to_first_token_seconds": [1.5]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": False, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == SLOW_TTFT_ID + + @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", @@ -318,3 +398,81 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la assert picked is not None assert picked["model_info"]["id"] == DEPLOYMENT_ID + + +def _latency_router(routing_strategy_args: dict) -> Router: + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in (FAST_TTFT_ID, SLOW_TTFT_ID) + ], + routing_strategy="latency-based-routing", + routing_strategy_args=routing_strategy_args, + ) + + +def _seed_streaming_ttft(router: Router) -> None: + router.cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": [0.1, 0.1, 1.0]}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": [0.3, 0.3, 0.3]}, + }, + ) + + +async def _pick_streaming(router: Router) -> str: + picked = await router.async_get_available_deployment( + model=MODEL_GROUP, + request_kwargs={"stream": True, "metadata": {}}, + ) + return picked["model_info"]["id"] + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_applies_ttft_percentile(): + """A config reload that adds ttft_percentile must reach the live selector, + not sit unused until the proxy restarts.""" + router = _latency_router({"max_latency_list_size": 50}) + _seed_streaming_ttft(router) + + assert await _pick_streaming(router) == SLOW_TTFT_ID + + router.update_settings(routing_strategy_args={"max_latency_list_size": 50, "ttft_percentile": 0.5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_keeps_previous_args_when_invalid(): + router = _latency_router({"ttft_percentile": 0.5}) + _seed_streaming_ttft(router) + + router.update_settings(routing_strategy_args={"ttft_percentile": 5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_is_a_noop_without_a_selector(): + """simple-shuffle has no selector to re-link, so an args update must leave + the router alone instead of blowing up on a missing selector attribute.""" + router = Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": FAST_TTFT_ID}, + } + ], + routing_strategy="simple-shuffle", + ) + + router.update_settings(routing_strategy_args={"ttl": 5}) + + assert router.routing_strategy_args == {"ttl": 5} + assert await _pick_streaming(router) == FAST_TTFT_ID From c7163a80ddfd34da743b06d770d60aba05b292eb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:48:06 -0700 Subject: [PATCH 108/157] perf(proxy): collapse per-worker SGR upserts into one statement per flush (#40362) Each proxy worker flushed one Prisma upsert per active (date, category, route) bucket every interval, so the Postgres primary saw workers x routes statements per interval across the deployment. A flush now builds a single multi-row INSERT ... ON CONFLICT DO UPDATE, and with use_redis_transaction_buffer on the workers push snapshots to a Redis list that one lease-holding pod folds and commits, so the whole deployment costs one statement per interval. The leader keeps popping until the list is empty so a deployment wider than the dequeue cap cannot build a backlog, and rows that fail both the commit and the Redis re-queue fall back to the leader's own accumulator instead of being lost. Resolves LIT-7371 Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/db/gateway_request_tracking.py | 220 ++++++++-- litellm/proxy/proxy_server.py | 14 +- .../proxy/db/test_gateway_request_tracking.py | 391 +++++++++++++++--- 4 files changed, 537 insertions(+), 89 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 108a914e9c1..f9389d22dea 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -335,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv( ########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer" +REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer" diff --git a/litellm/proxy/db/gateway_request_tracking.py b/litellm/proxy/db/gateway_request_tracking.py index bebd74e877c..c9ace68db33 100644 --- a/litellm/proxy/db/gateway_request_tracking.py +++ b/litellm/proxy/db/gateway_request_tracking.py @@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can add a key, so the fold and the table it commits to are bounded by (days x routes) however much traffic arrives, and the response path carries no unbounded queue that would block once full. + +A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO +UPDATE`` rather than one upsert per key, so a worker costs the primary one +statement per interval however many routes it served. With +``use_redis_transaction_buffer`` on, workers instead push their snapshot to a +Redis list and one lock-holding pod folds every entry and writes the table, so +the deployment as a whole costs the primary one statement per interval. """ -from dataclasses import asdict +import json +from collections.abc import AsyncIterator, Iterable from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory from litellm.types.proxy.gateway_requests import ( GatewayRequestCounts, @@ -28,6 +43,15 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient _EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0) +_TABLE: Final = '"LiteLLM_DailyGatewayRequests"' +_COLUMNS_PER_ROW: Final = 5 +_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')" +GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job" + +_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...] +_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows) +_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...]) +_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({}) def _utc_date() -> str: @@ -59,20 +83,54 @@ class GatewayRequestAccumulator: route) however long the database is unreachable. This buys at-least-once, not exactly-once, and the cost is worth stating. - The batch commits inside its context manager's ``__aexit__``, so a failure - raised after the transaction committed (a connection dropped while reading - the acknowledgement) restores counts that are already persisted, and the - next flush increments them a second time. Exactly-once would need a dedup - key the upserts could ignore on replay. For a traffic-volume metric a rare + The statement commits on the server before its acknowledgement is read, so + a failure raised after the commit (a connection dropped while reading the + acknowledgement) restores counts that are already persisted, and the next + flush increments them a second time. Exactly-once would need a dedup key + the upsert could ignore on replay. For a traffic-volume metric a rare overcount on a dropped acknowledgement beats losing a whole interval to every database blip, so the trade is deliberate. """ - for key, counts in snapshot.items(): - existing = self._counts.get(key, _EMPTY) - self._counts[key] = GatewayRequestCounts( - successful_requests=existing.successful_requests + counts.successful_requests, - failed_requests=existing.failed_requests + counts.failed_requests, - ) + self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced + + +def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot: + """Sum counts key-wise; the result stays bounded by (date x category x route).""" + folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once + for key, counts in items: + existing = folded.get(key, _EMPTY) + folded[key] = GatewayRequestCounts( + successful_requests=existing.successful_requests + counts.successful_requests, + failed_requests=existing.failed_requests + counts.failed_requests, + ) + return folded + + +def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]: + """ + One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category, + route) in the snapshot. Rows are ordered by the conflict key so concurrent + writers lock rows in the same order and cannot deadlock. + """ + ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + rows: Final = ", ".join( + f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})" + for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW) + ) + sql: Final = ( + f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n' + f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n' + f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n' + f' "updated_at" = {_UTC_NOW}' + ) + params: Final[tuple[str | int, ...]] = tuple( + value + for key, counts in ordered + for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + ) + return sql, params async def commit_gateway_requests_to_db( @@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db( prisma_client: "PrismaClient", snapshot: GatewayRequestSnapshot, ) -> None: - """Upsert one incrementing row per (date, category, route).""" + """Increment every (date, category, route) in the snapshot with a single statement.""" if not snapshot: return - ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + sql, params = build_gateway_requests_upsert(snapshot) + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client - # pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped, - # so .db and every table action off it resolve to Any at this boundary. The dict - # literals below are the shape prisma's generated inputs require. - async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client - for key, counts in ordered: - columns = asdict(key) - batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client - where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped - data={ # mutable-ok: prisma input is dict-shaped - "create": { # mutable-ok: prisma input is dict-shaped - **columns, - "successful_requests": counts.successful_requests, - "failed_requests": counts.failed_requests, - }, - "update": { # mutable-ok: prisma input is dict-shaped - "successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above - "failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above - }, - }, + verbose_proxy_logger.debug( + "Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot) + ) + + +class GatewayRequestRedisBuffer: + """ + Folds every worker's snapshot through one Redis list so a single pod per + interval writes the table, mirroring the spend writer's transaction buffer. + + Each entry is one worker's snapshot as JSON rows; the lock holder pops them, + sums them, and commits one statement. A commit failure pushes the summed + rows back so the next holder retries, keeping the at-least-once guarantee. + If that push fails too, the rows go back to the holder's own accumulator so + they ride along with its next flush instead of vanishing with the pop. + """ + + def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None: + self._redis_cache: Final = redis_cache + self._pod_lock_manager: Final = pod_lock_manager + + async def push(self, snapshot: GatewayRequestSnapshot) -> None: + if not snapshot: + return + rows: Final[_BufferedRows] = tuple( + (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + for key, counts in snapshot.items() + ) + await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),)) + + async def _pop_batch(self) -> tuple[str | bytes, ...]: + popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any + key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ) + if not popped: + return () + return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,)) + + async def _pop_all(self) -> AsyncIterator[str | bytes]: + while True: + batch = await self._pop_batch() + for entry in batch: + yield entry + if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT: + return + + async def pop(self) -> GatewayRequestSnapshot: + entries: Final = tuple([entry async for entry in self._pop_all()]) + return fold_counts( + ( + GatewayRequestKey(date=date, category=category, route=route), + GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed), ) + for entry in entries + for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry) + ) - verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered)) + async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot: + """ + Drain the list and write it as one statement, but only on the pod holding the job lock. + + The lock is a lease, never released: the holder re-enters it on every flush and + keeps committing alone until the TTL lapses, so the primary sees one statement + per flush interval deployment-wide instead of one per worker. + + Returns the popped rows that could be neither committed nor re-queued, for the + caller to keep in memory. Empty on success. + """ + if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME): + return _NO_COUNTS + buffered: Final = await self.pop() + try: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered) + except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler + verbose_proxy_logger.warning( + "Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush", + len(buffered), + exc_info=True, + ) + return await self._requeue(buffered) + return _NO_COUNTS + + async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot: + try: + await self.push(snapshot) + except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead + verbose_proxy_logger.warning( + "Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush", + len(snapshot), + exc_info=True, + ) + return snapshot + return _NO_COUNTS async def flush_gateway_requests( prisma_client: "PrismaClient", accumulator: GatewayRequestAccumulator, + redis_buffer: GatewayRequestRedisBuffer | None = None, ) -> None: """ Scheduler entrypoint. Never raises: a metering failure must not kill the job. + With ``redis_buffer`` the snapshot goes to Redis and only the lease holder + writes to Postgres. Shutdown passes no buffer so a departing worker writes its + own counts directly instead of parking them behind a lease it may not hold. + ``CancelledError`` is deliberately not caught, so a flush cancelled during shutdown drops its snapshot rather than restoring counts onto an accumulator the process is about to discard. """ snapshot: Final = accumulator.drain() try: - await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + if redis_buffer is None: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + else: + await redis_buffer.push(snapshot) except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler accumulator.restore(snapshot) verbose_proxy_logger.warning( @@ -131,3 +269,13 @@ async def flush_gateway_requests( len(snapshot), exc_info=True, ) + return + if redis_buffer is None: + return + try: + accumulator.restore(await redis_buffer.commit_if_leader(prisma_client)) + except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush + verbose_proxy_logger.warning( + "Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush", + exc_info=True, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c241e66049b..09e43eb74e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -424,6 +424,7 @@ from litellm.proxy.db.exception_handler import ( ) from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, + GatewayRequestRedisBuffer, flush_gateway_requests, ) from litellm.proxy.db.proxy_worker_heartbeat import ( @@ -2355,6 +2356,17 @@ open_telemetry_logger: OpenTelemetry | None = None gateway_request_accumulator: Final = GatewayRequestAccumulator() ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) + + +def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None: + """Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + redis_cache: Final = writer.redis_update_buffer.redis_cache + if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + return None + return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager) + + ### REDIS QUEUE ### async_result: Final = None celery_app_conn: Final = None @@ -9633,7 +9645,7 @@ class ProxyStartupEvent: flush_gateway_requests, "interval", seconds=batch_writing_interval, - args=(prisma_client, gateway_request_accumulator), + args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()), id="update_gateway_requests_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, diff --git a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py index 93a11a914cb..045261e2d53 100644 --- a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py +++ b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py @@ -8,8 +8,11 @@ from datetime import datetime, timezone import pytest +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY from litellm.proxy.db.gateway_request_tracking import ( + GATEWAY_REQUESTS_JOB_NAME, GatewayRequestAccumulator, + GatewayRequestRedisBuffer, commit_gateway_requests_to_db, flush_gateway_requests, ) @@ -83,40 +86,47 @@ def test_drain_snapshot_is_not_mutated_by_later_records(): # ── commit ──────────────────────────────────────────────────────────────────── -class FakeTable: - def __init__(self) -> None: - self.upserts: list[dict] = [] - - def upsert(self, *, where: dict, data: dict) -> None: - self.upserts.append({"where": where, "data": data}) - - -class FakeBatcher: - def __init__(self, table: FakeTable) -> None: - self.litellm_dailygatewayrequests = table - - async def __aenter__(self) -> "FakeBatcher": - return self - - async def __aexit__(self, *args: object) -> bool: - return False - - class FakeDB: - def __init__(self, table: FakeTable) -> None: - self._table = table + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] - def batch_(self) -> FakeBatcher: - return FakeBatcher(self._table) + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) // 5 class FakePrismaClient: def __init__(self) -> None: - self.table = FakeTable() - self.db = FakeDB(self.table) + self.db = FakeDB() -def test_commit_upserts_one_incrementing_row_per_key(): +def _rows_written(client: FakePrismaClient) -> list[tuple[object, ...]]: + """Every (date, category, route, successful, failed) tuple the database received, in statement order.""" + return [params[i : i + 5] for _, params in client.db.statements for i in range(0, len(params), 5)] + + +def test_commit_increments_with_a_single_statement_for_the_whole_snapshot(): + """One statement per flush is the whole point: the previous per-key upsert cost + the primary (workers x routes) statements per interval.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route=route): ( + GatewayRequestCounts(successful_requests=7, failed_requests=2) + ) + for route in ("/chat/completions", "/embeddings", "/responses", "/v1/messages", "/mcp") + } + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + assert len(client.db.statements) == 1 + sql, params = client.db.statements[0] + assert sql.count("ON CONFLICT") == 1 + assert sql.count("(NOW() AT TIME ZONE 'UTC'))") == 5 + assert len(params) == 25 + + +def test_commit_sql_adds_to_the_existing_row_instead_of_replacing_it(): + """A worker only knows its own share; the SQL must add EXCLUDED onto the stored count.""" client = FakePrismaClient() snapshot = { GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( @@ -126,20 +136,36 @@ def test_commit_upserts_one_incrementing_row_per_key(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - assert len(client.table.upserts) == 1 - written = client.table.upserts[0] - assert written["where"] == { - "date_category_route": { - "date": "2026-08-01", - "category": "llm", - "route": "/chat/completions", - } + sql, params = client.db.statements[0] + assert 'INSERT INTO "LiteLLM_DailyGatewayRequests"' in sql + assert 'ON CONFLICT ("date", "category", "route") DO UPDATE SET' in sql + assert ( + '"successful_requests" = "LiteLLM_DailyGatewayRequests"."successful_requests" + EXCLUDED."successful_requests"' + in sql + ) + assert '"failed_requests" = "LiteLLM_DailyGatewayRequests"."failed_requests" + EXCLUDED."failed_requests"' in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 7, 2) + + +def test_commit_placeholders_line_up_with_params(): + """$n positions are generated per row; a drift here silently swaps a route for a count.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=0) + ), + GatewayRequestKey(date="2026-08-01", category="mcp", route="/mcp"): ( + GatewayRequestCounts(successful_requests=0, failed_requests=3) + ), } - assert written["data"]["update"] == { - "successful_requests": {"increment": 7}, - "failed_requests": {"increment": 2}, - } - assert written["data"]["create"]["successful_requests"] == 7 + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + sql, params = client.db.statements[0] + assert "($1::text, $2::text, $3::text, $4::bigint, $5::bigint," in sql + assert "($6::text, $7::text, $8::text, $9::bigint, $10::bigint," in sql + assert "$11" not in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 1, 0, "2026-08-01", "mcp", "/mcp", 0, 3) def test_commit_is_deterministically_ordered(): @@ -154,17 +180,14 @@ def test_commit_is_deterministically_ordered(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - written_order = [ - (row["where"]["date_category_route"]["date"], row["where"]["date_category_route"]["category"]) - for row in client.table.upserts - ] + written_order = [(row[0], row[1]) for row in _rows_written(client)] assert written_order == [("2026-08-01", "llm"), ("2026-08-01", "mcp"), ("2026-08-02", "llm")] def test_commit_skips_the_database_entirely_when_nothing_accumulated(): client = FakePrismaClient() asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot={})) - assert client.table.upserts == [] + assert client.db.statements == [] # ── flush ───────────────────────────────────────────────────────────────────── @@ -177,12 +200,12 @@ def test_flush_drains_and_commits(): asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 + assert len(client.db.statements) == 1 assert acc.drain() == {} class ExplodingDB: - def batch_(self): + async def execute_raw(self, query: str, *args: object) -> int: raise RuntimeError("db gone") @@ -208,10 +231,7 @@ def test_failed_flush_keeps_counts_for_the_next_attempt(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert client.table.upserts[0]["data"]["update"] == { - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 1}, - } + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] def test_restored_counts_merge_with_requests_recorded_meanwhile(): @@ -223,5 +243,272 @@ def test_restored_counts_merge_with_requests_recorded_meanwhile(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 - assert client.table.upserts[0]["data"]["update"]["successful_requests"] == {"increment": 2} + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingDBWithInFlightRequest: + """Fails the write after a request has been recorded while it was in flight.""" + + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.accumulator = accumulator + + async def execute_raw(self, query: str, *args: object) -> int: + _record(self.accumulator, 500) + raise RuntimeError("db gone") + + +class ExplodingClientWithInFlightRequest: + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.db = ExplodingDBWithInFlightRequest(accumulator) + + +def test_restore_keeps_requests_recorded_while_the_failed_write_was_in_flight(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(ExplodingClientWithInFlightRequest(acc), acc)) + + client = FakePrismaClient() + asyncio.run(flush_gateway_requests(client, acc)) + + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] + + +# ── redis buffer ────────────────────────────────────────────────────────────── + + +class FakeRedis: + def __init__(self) -> None: + self.lists: dict[str, list[str]] = {} + + async def async_rpush(self, key: str, values: list[str]) -> int: + self.lists.setdefault(key, []).extend(values) + return len(self.lists[key]) + + async def async_lpop(self, key: str, count: int) -> list[str] | None: + queue = self.lists.get(key, []) + if not queue: + return None + popped, self.lists[key] = queue[:count], queue[count:] + return popped + + +class FakePodLock: + def __init__(self, *, leader: bool) -> None: + self.leader = leader + self.held: list[str] = [] + self.released: list[str] = [] + + async def acquire_lock(self, cronjob_id: str) -> bool: + self.held.append(cronjob_id) + return self.leader + + async def release_lock(self, cronjob_id: str) -> None: + self.released.append(cronjob_id) + + +class FakeLease: + """Redis-side view of the job lock: SET NX by pod id, re-entrant for the holder, freed only by release or TTL.""" + + def __init__(self) -> None: + self.holder: str | None = None + + +class FakeLeasePodLock: + def __init__(self, lease: FakeLease, pod_id: str) -> None: + self.lease = lease + self.pod_id = pod_id + + async def acquire_lock(self, cronjob_id: str) -> bool: + if self.lease.holder is None: + self.lease.holder = self.pod_id + return self.lease.holder == self.pod_id + + async def release_lock(self, cronjob_id: str) -> None: + if self.lease.holder == self.pod_id: + self.lease.holder = None + + +def _buffer(redis: FakeRedis, *, leader: bool) -> tuple[GatewayRequestRedisBuffer, FakePodLock]: + lock = FakePodLock(leader=leader) + return GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=lock), lock # pyright: ignore[reportArgumentType] # duck-typed fakes + + +def test_non_leader_workers_push_to_redis_and_never_touch_the_database(): + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(3): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + assert client.db.statements == [] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_folds_every_workers_snapshot_into_one_statement(): + """Fifty workers each flushing the same routes must cost the primary one statement, not fifty.""" + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(50): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500, route="/responses") + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader_acc = GatewayRequestAccumulator() + _record(leader_acc, 200) + leader, lock = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [ + (_today(), "llm", "/chat/completions", 51, 0), + (_today(), "llm", "/responses", 0, 50), + ] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + assert lock.held == [GATEWAY_REQUESTS_JOB_NAME] + assert lock.released == [] + + +def test_leader_keeps_the_lease_so_staggered_pods_cost_one_statement_per_interval(): + """Pods flush on their own clocks; without the lease each one would win the lock in turn and commit alone.""" + redis = FakeRedis() + client = FakePrismaClient() + lease = FakeLease() + pods = tuple( + GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=FakeLeasePodLock(lease, f"pod-{i}")) # pyright: ignore[reportArgumentType] # duck-typed fakes + for i in range(4) + ) + + for _interval in range(3): + for pod in pods: + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(client, acc, pod)) + + assert lease.holder == "pod-0" + assert len(client.db.statements) == 3 + assert [row[3] for row in _rows_written(client)] == [1, 4, 4] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_drains_a_backlog_deeper_than_one_capped_pop(): + """More workers than MAX_REDIS_BUFFER_DEQUEUE_COUNT must not leave a growing tail queued behind the cap.""" + redis = FakeRedis() + client = FakePrismaClient() + workers = MAX_REDIS_BUFFER_DEQUEUE_COUNT * 2 + 1 + for _ in range(workers): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", workers, 0)] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + +def test_leader_with_nothing_buffered_writes_nothing(): + redis = FakeRedis() + client = FakePrismaClient() + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert client.db.statements == [] + assert lock.released == [] + + +def test_leader_requeues_to_redis_when_the_database_commit_fails(): + """Counts popped from Redis are gone from every worker; a failed commit must put them back.""" + redis = FakeRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 200) + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(ExplodingClient(), acc, leader)) + + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + assert lock.released == [] + assert acc.drain() == {} + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingRedis(FakeRedis): + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis gone") + + +class UnreadableRedis(FakeRedis): + async def async_lpop(self, key: str, count: int) -> list[str] | None: + raise RuntimeError("redis gone mid-flush") + + +class UnwritableRedis(FakeRedis): + """Pops succeed, pushes fail: a Redis that went read-only between the leader's pop and its re-queue.""" + + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis read-only") + + +def test_leader_keeps_popped_counts_in_memory_when_both_the_database_and_the_requeue_fail(): + """The pop removed the only copy; if Redis will not take it back the leader itself must carry it.""" + redis = FakeRedis() + worker_acc = GatewayRequestAccumulator() + _record(worker_acc, 200) + _record(worker_acc, 200) + worker, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(FakePrismaClient(), worker_acc, worker)) + + degraded = UnwritableRedis() + degraded.lists = redis.lists + leader_acc = GatewayRequestAccumulator() + leader, _ = _buffer(degraded, leader=True) + asyncio.run(flush_gateway_requests(ExplodingClient(), leader_acc, leader)) + assert degraded.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +def test_leader_whose_redis_read_fails_leaves_the_pushed_rows_for_the_next_flush(): + """The scheduler job must not raise, and nothing is popped so nothing needs restoring anywhere.""" + redis = UnreadableRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + client = FakePrismaClient() + leader, _ = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, acc, leader)) + + assert client.db.statements == [] + assert acc.drain() == {} + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + + +def test_failed_redis_push_keeps_counts_locally_for_the_next_flush(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500) + buffer, lock = _buffer(ExplodingRedis(), leader=True) + + asyncio.run(flush_gateway_requests(FakePrismaClient(), acc, buffer)) + + assert lock.held == [] + assert acc.drain() == { + GatewayRequestKey(date=_today(), category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=1) + ) + } From 6195fbf5f3aaa656d916f0a7c3013ada34e87752 Mon Sep 17 00:00:00 2001 From: Animesh Kumar Date: Wed, 9 Sep 2026 23:24:12 +0530 Subject: [PATCH 109/157] test: isolate bedrock aws tests from ambient SSL env vars Nine cases assert the sts client is built with verify=True, but get_ssl_verify reads SSL_CERT_FILE and SSL_VERIFY, so the argument depended on the ambient environment. The published images set SSL_CERT_FILE, so the suite failed there while passing in CI. Fixes #40357 --- tests/test_litellm/llms/bedrock/test_base_aws_llm.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index f854d806bdc..a1c28f36e70 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -38,6 +38,14 @@ def flush_shared_bedrock_iam_cache(): yield +@pytest.fixture(autouse=True) +def _clean_ssl_env(monkeypatch): + """get_ssl_verify reads these, so the sts client's verify= would otherwise depend on + the ambient environment. The published images set SSL_CERT_FILE.""" + for env_var in ("SSL_CERT_FILE", "SSL_VERIFY"): + monkeypatch.delenv(env_var, raising=False) + + def test_base_aws_llm_instances_share_process_wide_iam_cache(): """Regression LIT-2662: new instances must reuse iam_cache (Bedrock passthrough is per-request).""" first = BaseAWSLLM() From c82c9cbcedf43661521270740273d6437d76c098 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 11:04:11 -0700 Subject: [PATCH 110/157] fix(router): strip encrypted reasoning on an auto-router tier change instead of a 503 (#40280) A Responses API follow-up that replays reasoning.encrypted_content is pinned to the deployment that minted it. Behind an auto-router the pre-routing hook rebinds the model to the tier it picked before the candidate pool is built, so a turn that classifies into a different tier never finds the origin and the affinity check raised its fail-fast 503, whose text claims a cooldown that does not exist When the deployment that minted the reasoning is not a member of the model group this turn is routed to, strip the encrypted reasoning (keeping any readable summary, string or block form) and dispatch to the routed group. Membership is tested by deployment id against the candidate set the router itself resolved for the route (routing group, model_name, team, and pattern alike), not by model-group name, so an alias, a provider-qualified spelling, a team-public name, or a pattern route of the same group is not misread as a tier change. An unknown origin (a removed deployment, or a forged/unauthenticated marker) is handled the same as a cross-group one and its reasoning is stripped, so a real cross-group id and a nonexistent id return the same response and cannot be used to enumerate deployment ids. Unavailability within the origin's own group keeps the existing 429/503 fail-fast, so the cooldown contract is unchanged Resolves LIT-7195 Claude-Session: https://claude.ai/code/session_01KAumQbhzk6jdWWHFLA8Jar Co-authored-by: Claude Opus 4.8 --- litellm/responses/utils.py | 43 ++ litellm/router.py | 34 ++ .../encrypted_content_affinity_check.py | 73 ++- .../test_encrypted_content_affinity_check.py | 561 ++++++++++++------ tests/test_litellm/test_router.py | 42 ++ 5 files changed, 557 insertions(+), 196 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 540d492beec..599e978df6a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -543,6 +543,49 @@ class ResponsesAPIRequestUtils: return request_input + @staticmethod + def strip_encrypted_reasoning_from_input(request_input: object) -> None: + """Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary. + + Mutates ``request_input`` in place: the router's fallback snapshot shares this + list object, so a rebound list would replay the stripped items on the fallback hop. + """ + if not isinstance(request_input, list): + return + items: Final = cast(list[object], request_input) # cast-ok: untyped client json + stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items) + items[:] = (item for item in stripped if item is not None) # rebind-ok: list shared with fallback snapshot + + @staticmethod + def _without_encrypted_reasoning(item: object) -> object | None: + if not isinstance(item, dict): + return item + reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json + if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): + return reasoning + readable: Final = any( + ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") + ) + if not readable: + return None + kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys + key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id") + } + return kept + + @staticmethod + def _has_readable_text(value: object) -> bool: + """A reasoning item's ``summary``/``content`` carries readable text: a non-empty string, or a + list holding at least one block with a non-empty ``text`` field (summary_text / output_text).""" + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, list): + return any( + isinstance(block, dict) and bool(cast(Mapping[str, object], block).get("text")) # cast-ok: untyped json + for block in value + ) + return False + @staticmethod def _build_responses_api_response_id( custom_llm_provider: str | None, diff --git a/litellm/router.py b/litellm/router.py index 2fbdf541487..68ace283949 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11167,6 +11167,40 @@ class Router: return ids + def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: + """ + Deployment ids that could serve ``model`` for ``team_id``, unioned across the paths + the router resolves a route through: ``model_group_alias``, a routing group, the + ``model_name`` and team indexes, and wildcard pattern routes. Read-only and + side-effect-free, unlike ``_common_checks_available_deployment`` which also applies + fallbacks and can raise. Lets a pre-call check tell a genuine cross-group route from + same-group unavailability without re-deriving that precedence at the call site, and + without leaking deployment ids into request kwargs bound for the provider. + """ + resolved: Final = self._get_model_from_alias(model=model) or model + routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) + if routing_group_members is not None: + return self._deployment_ids(routing_group_members) + if resolved in self.model_names: + return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id)) + team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None + return self._deployment_ids( + ( + *self._get_all_deployments(model_name=resolved, team_id=team_id), + *(self.pattern_router.route(resolved) or ()), + *((team_router.route(resolved) or ()) if team_router is not None else ()), + ) + ) + + @staticmethod + def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: + return frozenset( + str(model_info["id"]) + for deployment in deployments + for model_info in (deployment.get("model_info"),) + if isinstance(model_info, Mapping) and model_info.get("id") is not None + ) + def has_model_id(self, candidate_id: str) -> bool: """ O(1) membership check for a deployment ID without allocating large lists. diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..d10153881d9 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,13 +37,13 @@ Safe to enable globally: """ import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx from litellm._logging import verbose_router_logger from litellm.exceptions import ( - BadRequestError, RateLimitError, ServiceUnavailableError, ) @@ -158,6 +158,23 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None + @staticmethod + def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None: + containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping)) + return next((tid for tid in team_ids if isinstance(tid, str)), None) + + def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]: + """ + Deployment ids that could serve this turn's routed ``model``, as the router + resolves a route (model_group_alias / routing group / model_name / team / + pattern). Delegates to the router so the full precedence is not re-derived here + and no deployment ids are written into request kwargs bound for the provider. + """ + if self.router is None: + return frozenset() + return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs)) + @staticmethod def _encryption_boundary_key( litellm_params: object, @@ -225,10 +242,14 @@ class EncryptedContentAffinityCheck(CustomLogger): """ If the request ``input`` contains litellm-encoded item IDs, decode the embedded ``model_id`` and pin the request to that deployment. Raises - ``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError`` - when the originating deployment is unavailable and no encryption-boundary - peer exists, rather than dispatching a doomed request to a non-peer - deployment. The 429/503 split mirrors the originating cooldown's status: + ``RateLimitError`` / ``ServiceUnavailableError`` when the originating + deployment is a member of the routed model group but currently unavailable + and no encryption-boundary peer exists, rather than dispatching a doomed + request to a non-peer deployment. When the origin is not a member of the + routed group (an auto-router tier change, a model switch with no peer, a + removed deployment, or an unknown/forged marker), the encrypted reasoning is + stripped and the request dispatches with its readable history instead. The + 429/503 split mirrors the originating cooldown's status: a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the remaining cooldown window) so OpenAI-compatible clients back off and retry after the deployment is eligible again. @@ -285,12 +306,34 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - # Dispatching to a non-peer would guarantee an upstream - # `invalid_encrypted_content` 400, so fail fast with a clearer error. + # The origin cannot serve this turn's routed group and no peer shares the boundary, so its + # encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch + # to the routed group instead of failing. Membership is tested by deployment id against the set + # the router actually resolved for this route, not by model-group name, so an alias, a + # provider-qualified spelling, a team-public name, or a pattern route of the same group is not + # mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is + # treated the same as a cross-group one, which also denies an authenticated caller a + # deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and + # dispatch rather than returning distinguishable responses. Only a genuine same-group member + # that is currently unavailable falls through to the fail-fast, preserving the cooldown contract. + routed_group_model_ids: Final = ( + self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset() + ) + if str(model_id) not in routed_group_model_ids: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; " + "forwarding without its encrypted reasoning", + model_id, + model, + ) + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + return typed_healthy_deployments + + # The origin is a member of the routed group but currently unavailable (cooled down); fail fast + # rather than dispatching to a non-peer, which would guarantee an upstream 400. raise await self._unavailable_origin_error( model=model, model_id=model_id, - originating=originating, parent_otel_span=parent_otel_span, ) @@ -298,25 +341,11 @@ class EncryptedContentAffinityCheck(CustomLogger): self, model: str, model_id: str, - originating: Deployment | None, parent_otel_span: Span | None, ) -> Exception: # Public error messages intentionally omit the originating ``model_id`` so # an authenticated caller forging encrypted-content markers cannot use the # error surface to enumerate which deployment IDs exist on this router. - if originating is None: - return BadRequestError( - message=( - "The deployment that produced this encrypted_content is no " - "longer configured on this router, and no deployment on the " - "same encryption boundary is available. Re-issue the request " - "without the stale encrypted_content items, or restore the " - "originating deployment." - ), - model=model, - llm_provider="", - ) - cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span) if cooldown is not None and str(cooldown.get("status_code")) == "429": diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..3961c8d74c2 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -16,12 +16,10 @@ The mechanism works without any cache and supports two encoding strategies: """ import time -from typing import List, Optional from unittest.mock import AsyncMock, patch import pytest - import litellm from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -68,9 +66,7 @@ class TestEncryptedItemIdCodec: def test_roundtrip(self): model_id = "deployment-1" original_item_id = "rs_abc123def456" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) assert encoded.startswith("encitem_") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None @@ -81,9 +77,7 @@ class TestEncryptedItemIdCodec: """Decoding must succeed even if base64 padding (=) was stripped in transit.""" model_id = "gpt-5.1-codex-openai-2" original_item_id = "rs_0efb96cb222403210069a01d5d52588196a9dc394ffdb89d00" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) # Strip any trailing '=' to simulate what happens in transit stripped = encoded.rstrip("=") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(stripped) @@ -100,9 +94,7 @@ class TestEncryptedItemIdCodec: """item_id values containing ';' must survive the roundtrip.""" model_id = "deployment-1" original_item_id = "rs_part1;part2;part3" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None assert decoded["item_id"] == original_item_id @@ -118,11 +110,7 @@ class TestUpdateEncryptedContentItemIds: {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) # Plain message item untouched assert result["output"][0]["id"] == "msg_abc" # Reasoning item with encrypted_content gets encoded @@ -133,16 +121,8 @@ class TestUpdateEncryptedContentItemIds: assert decoded["item_id"] == "rs_xyz" def test_no_op_when_model_id_is_none(self): - response = { - "output": [ - {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"} - ] - } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, None - ) - ) + response = {"output": [{"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}]} + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, None) assert result["output"][0]["id"] == "rs_xyz" @@ -151,9 +131,7 @@ class TestEncryptedContentWrapping: """Test wrapping encrypted_content with model_id metadata.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original_encrypted_data" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_content @@ -170,9 +148,7 @@ class TestEncryptedContentWrapping: ( model_id, content, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - plain_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(plain_content) assert model_id is None assert content == plain_content @@ -189,11 +165,7 @@ class TestEncryptedContentWrapping: }, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) assert result["output"][0].get("encrypted_content") is None wrapped = result["output"][1]["encrypted_content"] assert wrapped.startswith("litellm_enc:") @@ -210,19 +182,13 @@ class TestRestoreEncryptedContentItemIds: def test_restores_encoded_ids(self): model_id = "deployment-1" original_id = "rs_encrypted_item_456" - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_id - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_id) request_input = [ {"type": "message", "id": "msg_abc123", "role": "assistant"}, {"type": "reasoning", "id": encoded_id, "encrypted_content": "secret"}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["id"] == "msg_abc123" assert restored[1]["id"] == original_id @@ -230,33 +196,21 @@ class TestRestoreEncryptedContentItemIds: """Test that wrapped encrypted_content is unwrapped before forwarding.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original" - wrapped_content = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) - ) + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) request_input = [ {"type": "reasoning", "encrypted_content": wrapped_content}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["encrypted_content"] == original_content def test_no_op_for_plain_string_input(self): - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - "Hello world" - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input("Hello world") assert result == "Hello world" def test_no_op_for_unencoded_ids(self): request_input = [{"type": "message", "id": "msg_plain"}] - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert result[0]["id"] == "msg_plain" @@ -283,9 +237,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): "id": "msg_abc123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], }, { "type": "reasoning", @@ -347,9 +299,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): # The response must have rewritten the encrypted item's ID to encoded form encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), ( + f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + ) # Verify the encoded ID decodes back to the correct deployment + original ID decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded_item_id) @@ -371,9 +323,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) @pytest.mark.asyncio @@ -478,9 +430,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): # Extract encoded item ID from the first response output encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), f"Expected encitem_... but got {encoded_item_id!r}" # Follow-up with the encoded item ID — should pin to same deployment second_response = await router.aresponses( @@ -628,17 +578,13 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): if hasattr(first_item, "encrypted_content") else first_item.get("encrypted_content") ) - assert wrapped_content.startswith( - "litellm_enc:" - ), f"Expected wrapped content but got {wrapped_content[:50]}..." + assert wrapped_content.startswith("litellm_enc:"), f"Expected wrapped content but got {wrapped_content[:50]}..." # Verify we can extract model_id from wrapped content ( extracted_model_id, _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - wrapped_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped_content) assert extracted_model_id == first_model_id # Second request: use wrapped encrypted_content WITHOUT an ID (Codex behavior) @@ -653,9 +599,9 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) def test_encrypted_content_wrapping_preserves_original_content(): @@ -664,13 +610,9 @@ def test_encrypted_content_wrapping_preserves_original_content(): This is critical for streaming responses where content must round-trip correctly. """ model_id = "test-deployment-1" - original_encrypted_content = ( - "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - ) + original_encrypted_content = "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_encrypted_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_encrypted_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_encrypted_content @@ -691,9 +633,7 @@ def test_encrypted_content_wrapping_with_multiple_semicolons(): model_id = "deployment-with-semicolons" original_content = "gAAAAAB;some;content;with;semicolons" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) ( extracted_model_id, @@ -764,9 +704,7 @@ async def test_encrypted_content_affinity_preserves_litellm_metadata_for_respons request_kwargs=request_kwargs, ) - assert ( - request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True - ) + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True assert request_kwargs["litellm_metadata"]["model_info"] == {"id": "dep-1"} @@ -777,9 +715,7 @@ def test_encrypted_content_wrapping_empty_string(): model_id = "test-deployment" original_content = "" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") @@ -1132,9 +1068,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): "api_key": "fake-azure-resource-key-a", } - pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key( - pydantic_params - ) + pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key(pydantic_params) plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None @@ -1161,18 +1095,8 @@ def test_boundary_key_rejects_non_dict_like_inputs(): for bad in (None, [], "not a dict", 42, object()): assert EncryptedContentAffinityCheck._encryption_boundary_key(bad) is None - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "", "api_key": "k"} - ) - is None - ) - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "https://x"} - ) - is None - ) + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "", "api_key": "k"}) is None + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "https://x"}) is None # --------------------------------------------------------------------------- @@ -1180,10 +1104,11 @@ def test_boundary_key_rejects_non_dict_like_inputs(): # --------------------------------------------------------------------------- -def _make_originating_mock(api_base: str, api_key: str): +def _make_originating_mock(api_base: str, api_key: str, model_name: str = "gpt-5.4"): from unittest.mock import MagicMock originating = MagicMock() + originating.model_name = model_name originating.litellm_params.model_dump.return_value = { "api_base": api_base, "api_key": api_key, @@ -1192,19 +1117,23 @@ def _make_originating_mock(api_base: str, api_key: str): def _make_router_mock_with_cooldown( - originating, cooldown_entries: Optional[List[tuple]] = None + originating, + cooldown_entries: list[tuple] | None = None, + routed_group_model_ids: list[str] | None = None, ): """ Build a MagicMock router whose ``cooldown_cache.async_get_active_cooldowns`` - returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown). + returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown), and + whose ``get_candidate_model_ids_for_route`` returns ``routed_group_model_ids`` + (the deployment ids the router resolves for the routed model; defaulting to ``[]`` + — origin absent from the routed group, i.e. a tier change). """ from unittest.mock import AsyncMock, MagicMock mock_router = MagicMock() mock_router.get_deployment.return_value = originating - mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock( - return_value=list(cooldown_entries or []) - ) + mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock(return_value=list(cooldown_entries or [])) + mock_router.get_candidate_model_ids_for_route.return_value = frozenset(routed_group_model_ids or []) return mock_router @@ -1235,15 +1164,15 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42 }, ) ], + routed_group_model_ids=["deployment-a-cooled", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1297,15 +1226,15 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo }, ) ], + routed_group_model_ids=["deployment-a-cooled-429", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled-429", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled-429", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1345,15 +1274,16 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ ) originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") - mock_router = _make_router_mock_with_cooldown(originating, cooldown_entries=[]) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a-filtered", "deployment-b"] + ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-filtered", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-filtered", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1377,15 +1307,18 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ @pytest.mark.asyncio -async def test_affinity_raises_bad_request_when_origin_removed(): +async def test_affinity_strips_and_dispatches_when_origin_is_unknown_or_removed(): """ - Originating deployment was removed from the router config and no boundary - peer is available. This is permanent (the stale encrypted_content cannot - be honored), so surface a 400 with actionable text. + A removed deployment, or a forged/unknown affinity marker, resolves to no + originating deployment. It is handled like a cross-group origin: the encrypted + reasoning is stripped and the request dispatches with its readable history, + rather than returning a distinguishable error. That uniform handling denies an + authenticated caller a deployment-id existence oracle, an existing cross-group id + and a nonexistent id both strip and proceed, so responses cannot be told apart. + The membership lookup is skipped entirely when the origin is unknown. """ from unittest.mock import MagicMock - from litellm.exceptions import BadRequestError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1394,12 +1327,11 @@ async def test_affinity_raises_bad_request_when_origin_removed(): mock_router.get_deployment.return_value = None check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-removed", "rs_test" - ) - healthy_only_b = [ + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-removed") + routed_pool = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1408,18 +1340,28 @@ async def test_affinity_raises_bad_request_when_origin_removed(): } ] request_kwargs = { - "input": [{"id": encoded_id, "type": "reasoning"}], + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], } - with pytest.raises(BadRequestError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert "deployment-removed" not in str(excinfo.value) + assert result is routed_pool + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) + mock_router.get_candidate_model_ids_for_route.assert_not_called() @pytest.mark.asyncio @@ -1444,9 +1386,7 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): mock_router.get_deployment.return_value = originating check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_test") peer = { "model_info": {"id": "deployment-a-peer"}, "litellm_params": { @@ -1490,9 +1430,7 @@ async def test_model_group_affinity_config_enables_encrypted_content_affinity(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1536,9 +1474,7 @@ async def test_model_group_affinity_config_does_not_disable_global_encrypted_con }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1600,15 +1536,9 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen try: callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) assert encrypted_content_callback.enable_global_affinity is False cache_key = DeploymentAffinityCheck.get_affinity_cache_key( @@ -1620,9 +1550,7 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [ { @@ -1643,16 +1571,301 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen ) assert after_deployment_affinity == [deployment_a, deployment_b] - after_encrypted_content_affinity = ( - await encrypted_content_callback.async_filter_deployments( - model=model_group, - healthy_deployments=after_deployment_affinity, - messages=None, - request_kwargs=request_kwargs, - ) + after_encrypted_content_affinity = await encrypted_content_callback.async_filter_deployments( + model=model_group, + healthy_deployments=after_deployment_affinity, + messages=None, + request_kwargs=request_kwargs, ) assert after_encrypted_content_affinity == [deployment_b] assert request_kwargs.get("_encrypted_content_affinity_pinned") is True finally: router.discard() + + +class TestStripEncryptedReasoningFromInput: + def test_keeps_summary_and_drops_encrypted_content_and_id(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_1") + request_input = [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "id": encoded_id, + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "reasoning", "id": encoded_id, "encrypted_content": wrapped}, + {"type": "reasoning", "encrypted_content": wrapped, "summary": []}, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + + def test_keeps_string_form_summary_when_stripping(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + request_input = [ + {"type": "reasoning", "encrypted_content": wrapped, "summary": "plain string thought"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "content": [{"type": "output_text", "text": "in content"}], + }, + {"type": "reasoning", "encrypted_content": wrapped, "summary": "", "content": []}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"type": "reasoning", "summary": "plain string thought"}, + {"type": "reasoning", "content": [{"type": "output_text", "text": "in content"}]}, + ] + + def test_leaves_input_untouched_when_no_encrypted_reasoning(self): + request_input = [ + {"role": "user", "content": "first turn"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "no blob here"}]}, + {"role": "user", "content": "second turn"}, + ] + before = [dict(item) for item in request_input] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == before + + +def _cross_group_request_kwargs(): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + return { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "ZEBRA: why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"type": "message", "role": "assistant", "content": "Rayleigh scattering."}, + {"role": "user", "content": "KIWI: and sunsets?"}, + ], + } + + +@pytest.mark.asyncio +async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_group(): + """ + An auto-router tier change (or a model switch with no boundary peer): the + routed pool holds no deployment of the origin's model group. The origin is + healthy, so a 503 would be wrong; the follow-up dispatches to the routed + pool with the origin's encrypted reasoning stripped and its summary kept. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [ + { + "model_info": {"id": "deployment-b"}, + "model_name": "gpt-simple-tier", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5-nano", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + original_input = request_kwargs["input"] + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + assert request_kwargs["input"] is original_input + assert [item.get("type") or item["role"] for item in original_input] == [ + "user", + "reasoning", + "message", + "user", + ] + assert original_input[1] == { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "scattering"}], + } + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in original_input) + + +@pytest.mark.asyncio +async def test_affinity_fails_fast_within_the_origins_own_group(): + """ + Negative class for the tier-change discriminator: the routed group IS the + origin's group (a same-group cooldown, not a tier change), so even with a + healthy non-origin sibling that cannot decrypt the content, the request + still fails fast and the encrypted reasoning is left intact rather than + stripped. Preserves the LIT-3051 cooldown contract. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock( + "https://account-a.openai.azure.com/", "key-a", model_name="gpt-reasoning-tier" + ) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a", "deployment-a2"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + sibling_pool = [ + { + "model_info": {"id": "deployment-a2"}, + "model_name": "gpt-reasoning-tier", + "litellm_params": { + "api_base": "https://account-a2.openai.azure.com/", + "api_key": "key-a2", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-reasoning-tier", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id(): + """ + The discriminator must key on deployment-id membership, not on the model-group + name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini`` + while the routed group is the canonical ``gpt-5.4-mini``: same group, different + spelling. A name compare (``originating.model_name != model``) would read this as + a tier change and strip the reasoning it did not have to. Because the origin's id + is a member of the routed group, this is a same-group cooldown instead: the request + fails fast and the encrypted reasoning is left intact. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="openai/gpt-5.4-mini") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-mini-a", "deployment-mini-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-mini-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-mini-b"}, + "model_name": "gpt-5.4-mini", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-5.4-mini", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(): + """ + The exact `model_name` index does not include team-public or pattern routes, so a + same-group cooldown reached only through one of those would be misread as a tier change + and stripped. The check asks the router for the candidate ids it resolves for the route + (`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare + index. Here that set marks the origin as a candidate, so the request fails fast with its + reasoning intact, and the routed group and team are passed through to the router. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="model_name_teamA_uuid") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-team-a", "deployment-team-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-team-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-team-b"}, + "model_name": "team-public-model", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {"user_api_key_team_id": "teamA"}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="team-public-model", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA") diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bc79c5f6589..80da922724d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -14728,3 +14728,45 @@ def test_router_stays_quiet_when_a_deployment_drop_params_is_a_flag(value, caplo ) assert "is not a flag value" not in caplog.text + + +def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): + """ + get_candidate_model_ids_for_route resolves a route the way the router does, so a + pre-call check can tell a genuine cross-group route from same-group unavailability. + A concrete model group returns its member ids; a wildcard/pattern deployment is + included for a concrete model it matches, which the bare model_name index misses. + Regression guard for the LIT-7195 tier-change discriminator's team/pattern gaps. + """ + router = Router( + model_list=[ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-b", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-b"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-c", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-wild"}, + }, + ] + ) + + assert router.get_candidate_model_ids_for_route(model="grp") == frozenset({"dep-a", "dep-b"}) + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="openai/gpt-4o-some-new-model") + + +def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): + deployments = ( + {"model_info": {"id": "a"}}, + {"model_info": {"id": 2}}, + {"model_info": {}}, + {"no_model_info": True}, + ) + assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) From 096984bfc2f146e6e4bb6c259e8daab885a5e5f6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:19:29 -0700 Subject: [PATCH 111/157] fix(proxy): pin multi-root CA bundle to the server's root before handing it to Prisma (#40428) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_url_settings.py | 107 +++++++++- .../proxy/db/test_db_url_settings.py | 199 ++++++++++++++++-- 2 files changed, 282 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 4a0231ad9df..d1e4b3e92b8 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -34,12 +34,20 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the ones the reader URL does not pin itself. """ +import _ssl +import hashlib import os +import socket +import ssl +import struct +import sys +import tempfile import urllib.parse -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from functools import partial +from pathlib import Path from types import MappingProxyType -from typing import Annotated, Final, cast +from typing import Annotated, Final, Protocol, TypeAlias, cast from pydantic import AliasChoices, BeforeValidator, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -126,21 +134,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"}) +PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----" +PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103) +TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0 + +RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax -def translate_libpq_ssl_params(url: str) -> str: +class _VerifiedChainSource(Protocol): + def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ... + + +def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]: + if sys.version_info >= (3, 13): + return tuple(tls.get_verified_chain()) + legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10 + "_VerifiedChainSource | None", + tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13 + ) + chain: Final = () if legacy is None else legacy.get_verified_chain() or () + return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain) + + +def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None: + try: + context: Final = ssl.create_default_context(cafile=cafile) + with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw: + raw.sendall(PG_SSL_REQUEST) + if raw.recv(1) != b"S": + return None + with context.wrap_socket(raw, server_hostname=host) as tls: + chain: Final = _verified_chain_der(tls) + except (OSError, ValueError): + return None + return chain[-1] if chain else None + + +def pin_bundle_root(cert_path: str, host: str, port: int) -> str: + """Reduce a multi-root CA bundle to the one root that verifies ``host``. + + Prisma's ``sslcert`` loads a single PEM certificate (native-tls + ``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS + global bundle trusts only the first of its 108 regional roots and the + handshake fails with "unable to get local issuer certificate" for every + other region. A single-certificate file is returned as is. For a bundle, + one verifying handshake (chain and hostname, whole bundle as trust store) + identifies the trust anchor the server actually chains to, which is + written to a single-certificate file for Prisma. If the probe fails the + bundle path is returned unchanged, so Prisma fails closed exactly as + before rather than trusting anything the bundle would not. + """ + try: + if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2: + return cert_path + except OSError: + return cert_path + root: Final = _server_trust_anchor(cert_path, host, port) + if root is None: + return cert_path + pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem" + return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path + + +def _replace_file(target: Path, content: str) -> bool: + """Write ``content`` to a private temp file and rename it over ``target``, so + readers never see a partial file and a symlink planted at ``target`` is + replaced rather than followed.""" + try: + fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.") + except OSError: + return False + try: + with os.fdopen(fd, "w") as handle: + handle.write(content) + os.replace(staged, target) + except OSError: + Path(staged).unlink(missing_ok=True) + return False + return True + + +def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str: """Rewrite libpq's certificate-verification params into Prisma's dialect. Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert`` - (the CA bundle) and ``sslaccept=strict``. It silently discards + (a single CA certificate) and ``sslaccept=strict``. It silently discards ``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to ``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no certificate check at all. ``verify-ca`` and ``verify-full`` both become ``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes - ``sslcert``, and either one turns on ``sslaccept=strict`` (chain and - hostname), matching libpq where a root cert makes ``require`` verify. - Prisma params the operator pinned themselves win; anything else is left - untouched. + ``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root + bundle down to the server's root), and either one turns on + ``sslaccept=strict`` (chain and hostname), matching libpq where a root + cert makes ``require`` verify. Prisma params the operator pinned + themselves win; anything else is left untouched. """ parsed: Final = urllib.parse.urlsplit(url) pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) @@ -154,7 +241,9 @@ def translate_libpq_ssl_params(url: str) -> str: if key != "sslrootcert" ) root_cert: Final = tuple( - ("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys + ("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT))) + for key, value in pairs + if key == "sslrootcert" and "sslcert" not in keys ) strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),) query: Final = urllib.parse.urlencode(translated + root_cert + strict) diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index ba342342366..875dca4bee3 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -11,15 +11,30 @@ clobber a pre-existing ``DATABASE_URL_READ_REPLICA``. A pre-existing ``DATABASE_URL`` (password auth) is likewise left untouched. """ +import datetime +import hashlib import os +import socket +import ssl +import tempfile +import threading import urllib.parse +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final from unittest.mock import patch import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( + PG_SSL_REQUEST, DatabaseURLSettings, + translate_libpq_ssl_params, unsupported_db_scheme, unsupported_db_scheme_message, ) @@ -381,9 +396,7 @@ def test_writer_password_is_percent_encoded(monkeypatch): def test_writer_url_not_clobbered_when_already_set(monkeypatch): """An operator-pinned DATABASE_URL (e.g. helm's $(VAR) assembly) always wins over the discrete fields.""" - monkeypatch.setenv( - "DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db" - ) + monkeypatch.setenv("DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db") monkeypatch.setenv("DATABASE_HOST", "writer.example.com") monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm_db") @@ -515,9 +528,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db") with pytest.raises(RuntimeError, match=r"DATABASE_URL_READ_REPLICA.*mysql"): _apply() @@ -542,15 +553,11 @@ def test_reader_inherits_writer_connection_params(monkeypatch): "DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20&pgbouncer=true", ) - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db") _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["3"] assert query["pool_timeout"] == ["20"] assert query["pgbouncer"] == ["true"] @@ -568,9 +575,7 @@ def test_reader_keeps_its_own_pinned_connection_params(monkeypatch): _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["50"] assert query["pool_timeout"] == ["20"] @@ -776,6 +781,170 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa } +def _issue_cert( + subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool +) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name((x509.NameAttribute(x509.NameOID.COMMON_NAME, subject),)) + now: Final = datetime.datetime.now(datetime.timezone.utc) + builder: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(issuer.subject if issuer else name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=5)) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=ca, path_length=None), critical=True) + .add_extension(x509.SubjectAlternativeName((x509.DNSName("localhost"),)), critical=False) + ) + return builder.sign(issuer_key or key, hashes.SHA256()), key + + +def _pem(cert: x509.Certificate) -> bytes: + return cert.public_bytes(serialization.Encoding.PEM) + + +class _TlsPostgresStub: + """Answers one libpq ``SSLRequest`` with ``S`` and serves ``leaf + intermediate``.""" + + def __init__(self, chain_pem: Path, key_pem: Path) -> None: + self.context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + self.context.load_cert_chain(str(chain_pem), str(key_pem)) + self.listener: Final = socket.create_server(("127.0.0.1", 0)) + self.port: Final[int] = self.listener.getsockname()[1] + self.thread: Final = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + def _serve(self) -> None: + with self.listener: + while True: + try: + conn: socket.socket = self.listener.accept()[0] + except OSError: + return + with conn: + try: + if conn.recv(8) == PG_SSL_REQUEST: + conn.sendall(b"S") + with self.context.wrap_socket(conn, server_side=True) as tls: + tls.recv(1) + except OSError: + continue + + +@dataclass(frozen=True, slots=True) +class _RdsLikePki: + bundle: Path + wrong_bundle: Path + root: Path + port: int + + +@pytest.fixture +def rds_like_pki(tmp_path: Path) -> Iterator[_RdsLikePki]: + """An RDS-shaped trust setup: the server sends leaf + intermediate, the + bundle holds only self-signed roots, and the right root is not first.""" + root, root_key = _issue_cert("Real Root CA", None, None, ca=True) + decoys: Final = tuple(_issue_cert(f"Decoy Root CA {i}", None, None, ca=True)[0] for i in range(3)) + intermediate, intermediate_key = _issue_cert("Intermediate CA", root, root_key, ca=True) + leaf, leaf_key = _issue_cert("localhost", intermediate, intermediate_key, ca=False) + chain_pem: Final = tmp_path / "server-chain.pem" + chain_pem.write_bytes(_pem(leaf) + _pem(intermediate)) + key_pem: Final = tmp_path / "server.key" + key_pem.write_bytes( + leaf_key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + ) + bundle: Final = tmp_path / "global-bundle.pem" + bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys) + _pem(root)) + wrong_bundle: Final = tmp_path / "wrong-bundle.pem" + wrong_bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys)) + root_pem: Final = tmp_path / "root.pem" + root_pem.write_bytes(_pem(root)) + stub: Final = _TlsPostgresStub(chain_pem, key_pem) + yield _RdsLikePki(bundle=bundle, wrong_bundle=wrong_bundle, root=root_pem, port=stub.port) + stub.listener.close() + + +def _params(url: str) -> tuple[tuple[str, str], ...]: + return tuple(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True)) + + +def test_multi_root_bundle_is_pinned_to_the_root_the_server_chains_to( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Prisma's ``sslcert`` loads only the first certificate of the file, so + handing it the whole RDS bundle trusts one region's root and fails with + "unable to get local issuer certificate" everywhere else. The URL Prisma + receives must point at a single-certificate file holding the server's root.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + (sslmode, sslcert, sslaccept, _) = _params(os.environ["DATABASE_URL"]) + assert (sslmode, sslaccept) == (("sslmode", "require"), ("sslaccept", "strict")) + assert sslcert[0] == "sslcert" and sslcert[1] != str(rds_like_pki.bundle) + assert Path(sslcert[1]).read_bytes() == rds_like_pki.root.read_bytes() + + +def test_pinned_root_replaces_a_planted_symlink_instead_of_writing_through_it( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki, tmp_path: Path +): + """The pinned file has a predictable name in a shared temp dir, so a symlink + planted there must not redirect the write onto its target.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + root_der: Final = x509.load_pem_x509_certificate(rds_like_pki.root.read_bytes()).public_bytes( + serialization.Encoding.DER + ) + pinned: Final = tmp_path / f"litellm-sslcert-{hashlib.sha256(root_der).hexdigest()[:16]}.pem" + victim: Final = tmp_path / "victim.txt" + victim.write_text("untouched") + pinned.symlink_to(victim) + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + assert ("sslcert", str(pinned)) in _params(os.environ["DATABASE_URL"]) + assert victim.read_text() == "untouched" + assert not pinned.is_symlink() and pinned.read_bytes() == rds_like_pki.root.read_bytes() + + +def test_bundle_without_the_servers_root_is_passed_through_unchanged( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Nothing in the bundle verifies the server, so no root is pinned and + Prisma keeps rejecting the connection instead of trusting a root the + operator never shipped.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db" + f"?sslmode=verify-full&sslrootcert={rds_like_pki.wrong_bundle}", + ) + + _apply() + + assert ("sslcert", str(rds_like_pki.wrong_bundle)) in _params(os.environ["DATABASE_URL"]) + + +def test_root_cert_resolver_receives_the_urls_host_and_default_port(): + def resolver(cert_path: str, host: str, port: int) -> str: + return f"/pinned/{host}/{port}{cert_path}" + + url: Final = translate_libpq_ssl_params( + "postgresql://u:p@db.example.com/litellm_db?sslmode=verify-full&sslrootcert=/certs/bundle.pem", resolver + ) + + assert ("sslcert", "/pinned/db.example.com/5432/certs/bundle.pem") in _params(url) + + def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-ca") From 00e4380cb5b84b9681cbf6e14886cdbaf328cf4b Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 19:15:44 +0000 Subject: [PATCH 112/157] fix(registry): add Bedrock gpt-6-astra CRIS + mantle profiles, embed-v4/pegasus global profiles, gpt-image-2.5 entries; cap Vertex grok-4.1-fast output at 128k Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 188 +++++++++++++++++- model_prices_and_context_window.json | 188 +++++++++++++++++- 2 files changed, 368 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cdd9fc7b6d4..664c2737157 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -759,6 +759,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -14431,6 +14439,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -29748,6 +29778,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -48040,8 +48130,8 @@ "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 128000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -48056,8 +48146,8 @@ "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 128000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -55429,6 +55519,96 @@ "supports_reasoning": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cdd9fc7b6d4..664c2737157 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -759,6 +759,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -14431,6 +14439,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -29748,6 +29778,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -48040,8 +48130,8 @@ "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 128000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -48056,8 +48146,8 @@ "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 128000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -55429,6 +55519,96 @@ "supports_reasoning": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, From 34d2d010c3f0f08ab9c13700c7004d98377309fa Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 9 Sep 2026 19:33:59 +0000 Subject: [PATCH 113/157] test(cli): drop lite e2e tests, the e2e runner does not install the package Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/coverage_registry/other.yaml | 2 - tests/e2e/other/test_cli_cost_map_e2e.py | 104 ----------------------- 2 files changed, 106 deletions(-) delete mode 100644 tests/e2e/other/test_cli_cost_map_e2e.py diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index a7ba0dbb8cb..814ebae2e0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -48,5 +48,3 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} -- {id: other.cli.model_cost_map.version_skips_fetch, module: other, tier: P1, area: cli, assertions: [version_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite version command succeeds without requesting the remote model-cost map"} -- {id: other.cli.model_cost_map.models_list_skips_fetch, module: other, tier: P1, area: cli, assertions: [models_list_skips_fetch], source: "get_model_cost_map.py _is_cli_process / LIT-7385", fail_before_fix: proven, rationale: "The lite models list command uses the proxy without requesting the remote model-cost map"} diff --git a/tests/e2e/other/test_cli_cost_map_e2e.py b/tests/e2e/other/test_cli_cost_map_e2e.py deleted file mode 100644 index 62cd3d6c3f5..00000000000 --- a/tests/e2e/other/test_cli_cost_map_e2e.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations - -import os -import shutil -import subprocess -import threading -from collections.abc import Mapping -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Final - -import pytest -from e2e_config import MASTER_KEY, PROXY_BASE_URL -from proxy_client import ProxyClient - -pytestmark = pytest.mark.e2e - - -def _start_cost_map_server(request_log: Path) -> tuple[ThreadingHTTPServer, threading.Thread]: - class CostMapHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - with request_log.open("a", encoding="utf-8") as log_file: - log_file.write(f"{self.path}\n") - body: Final = b'{"test-model": {"litellm_provider": "openai", "mode": "chat"}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, format: str, *args: object) -> None: - return - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), CostMapHandler) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - return server, thread - - -def _lite_env(server: ThreadingHTTPServer, api_key: str | None) -> dict[str, str]: - base_env: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - return { - **base_env, - "LITELLM_MODEL_COST_MAP_URL": f"http://127.0.0.1:{server.server_port}/map.json", - "LITELLM_PROXY_URL": PROXY_BASE_URL, - **({"LITELLM_PROXY_API_KEY": api_key} if api_key is not None else {}), - } - - -def _run_lite( - args: tuple[str, ...], - server: ThreadingHTTPServer, - env: Mapping[str, str], -) -> subprocess.CompletedProcess[str]: - lite_path: Final = shutil.which("lite") - assert lite_path is not None, "the installed lite executable is required for e2e coverage" - try: - return subprocess.run( - [lite_path, *args], - capture_output=True, - text=True, - timeout=60, - env=env, - ) - finally: - server.shutdown() - server.server_close() - - -def _request_count(request_log: Path) -> int: - return request_log.read_text(encoding="utf-8").count("\n") if request_log.exists() else 0 - - -class TestLiteCliCostMapFetch: - @pytest.mark.covers("other.cli.model_cost_map.version_skips_fetch") - def test_lite_version_makes_no_cost_map_request(self, tmp_path: Path) -> None: - request_log: Final = tmp_path / "requests.log" - server, thread = _start_cost_map_server(request_log) - env: Final = _lite_env(server, None) - try: - result: Final = _run_lite(("--version",), server, env) - finally: - thread.join(timeout=10) - - assert result.returncode == 0 - assert "LiteLLM Proxy CLI Version" in result.stdout - assert _request_count(request_log) == 0 - - @pytest.mark.covers("other.cli.model_cost_map.models_list_skips_fetch") - def test_lite_models_list_uses_proxy_not_cost_map(self, tmp_path: Path, proxy: ProxyClient) -> None: - model_names: Final = tuple(entry.model_name for entry in proxy.model_info()) - assert model_names - request_log: Final = tmp_path / "requests.log" - server, thread = _start_cost_map_server(request_log) - env: Final = _lite_env(server, MASTER_KEY) - try: - result: Final = _run_lite(("models", "list"), server, env) - finally: - thread.join(timeout=10) - - assert result.returncode == 0 - assert result.stdout.strip() - assert any(model_name in result.stdout for model_name in model_names) - assert _request_count(request_log) == 0 From 8bb6d8c120e2c116a719495ecf2401a2bab9dd5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:34:08 -0700 Subject: [PATCH 114/157] chore(lint): bring ANN202 and BLE001 back under the strict-rule budget Annotate the return types of dispatch_async and transform_then_dispatch in llm_http_handler and _send_batch in azure_sentinel, and mark four legitimate broad catches with the repo's noqa convention, so the promote PR's lint job passes the strict gate again. Supersedes #40328. --- litellm/integrations/azure_sentinel/azure_sentinel.py | 4 +++- litellm/litellm_core_utils/get_model_cost_map.py | 2 +- litellm/litellm_core_utils/llm_cost_calc/utils.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 4 ++-- litellm/proxy/_experimental/mcp_server/server.py | 2 +- litellm/proxy/management_endpoints/cost_tracking_settings.py | 2 +- 6 files changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index eba6c862f7a..db5f790615f 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import Final, TypeVar from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( BatchSendCancelled, @@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger): "Content-Type": "application/json", } - async def _send_batch(batch: Sequence[_QueuedPayload]): + async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response: body: Final = safe_dumps(batch) return await self.async_httpx_client.post( url=api_endpoint, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index f81ddbfee2e..37375f25e46 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -585,7 +585,7 @@ def _retry_remote_fetch_in_background( _cost_map_source_info.fallback_reason = None _cost_map_source_info.loaded_at = datetime.now(timezone.utc) adopt_model_cost_map(finalized) - except Exception as e: + except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 5675c59733d..e5977ca4156 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1444,7 +1444,7 @@ def get_billed_token_rates( return _custom_pricing_rates(custom_cost_per_token) try: model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates return None return _cost_map_billed_rates( model_info=model_info, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7587b963a38..84b03e1955a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -579,7 +579,7 @@ class BaseLLMHTTPHandler: data: dict[str, object], # mutable-ok: async_completion takes dict signed_headers: dict[str, object], # mutable-ok: async_completion takes dict signed_json_body: bytes | None, - ): + ) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]: async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None if stream is True: return self.acompletion_stream_function( @@ -626,7 +626,7 @@ class BaseLLMHTTPHandler: if acompletion is True and provider_config.uses_async_transform_request: - async def transform_then_dispatch(): + async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict "dict[str, object]", await provider_config.async_transform_request( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a52da0cab1..f8f1c563811 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1053,7 +1053,7 @@ if MCP_AVAILABLE: route="/mcp/call_tool", traceback_str=failure_traceback, ) - except Exception: + except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error verbose_logger.exception("Error logging failed MCP proxy tool call") raise if proxy_logging_obj is not None: diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 493f75008c3..dc0da63555f 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -643,7 +643,7 @@ async def estimate_cost( custom_cost_per_token=resolved.custom_cost_per_token, litellm_logging_obj=litellm_logging_obj, ) - except Exception as e: + except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model raise HTTPException( status_code=404, detail={ From 35def27e7b145c7aef7463bef34db4e56a2e8d9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:37:29 -0700 Subject: [PATCH 115/157] fix(convert_dict_to_response): accept only a real list as choices and keep /v1/messages alive on an empty one Narrows the no-choices guard so a dict, string, or None still raises the APIError while an empty list passes through, guards the non-stream Anthropic bridge against indexing an empty choices list, and repairs test_completion_missing_role, whose raw-response mock was patched in as the create() callable itself so the handler only ever saw a MagicMock --- .../convert_dict_to_response.py | 8 ++--- .../adapters/transformation.py | 3 +- .../test_convert_dict_to_response.py | 30 +++++++++++++++++++ ...al_pass_through_adapters_transformation.py | 15 ++++++++++ tests/test_litellm/test_main.py | 2 +- 5 files changed, 52 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 7d8072ba622..29b75812bfb 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,7 @@ import json import re import time import traceback -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Final, Literal, cast import litellm @@ -179,7 +179,7 @@ async def convert_to_streaming_response_async( choice_list: Final[list[StreamingChoices]] = [] - if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( @@ -287,7 +287,7 @@ def convert_to_streaming_response( model_response_object: Final = ModelResponseStream() choice_list: Final[list[StreamingChoices]] = [] - if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( @@ -623,7 +623,7 @@ def convert_to_model_response_object( return convert_to_streaming_response(response_object=response_object) choice_list: Final[list[Choices]] = [] - if "choices" not in response_object or not isinstance(response_object["choices"], Iterable): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index db890662132..4bce760e943 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1487,8 +1487,9 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_content.insert(0, polyfill_result.compaction_block) ## extract finish reason + openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop" translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic( - openai_finish_reason=response.choices[0].finish_reason + openai_finish_reason=openai_finish_reason ) anthropic_finish_reason: Final = ( "refusal" diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 5406df691aa..ea50bff462e 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,3 +1,5 @@ +from typing import Final + import pytest from litellm.constants import RESPONSE_FORMAT_TOOL_NAME @@ -167,3 +169,31 @@ def test_convert_missing_choices_raises_api_error() -> None: ) assert "no 'choices'" in str(exc_info.value) + +@pytest.mark.parametrize("choices", [{}, "", None, 0]) +@pytest.mark.asyncio +async def test_convert_non_list_choices_raises_api_error(choices: object) -> None: + from litellm.exceptions import APIError + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": choices, + } + with pytest.raises(APIError, match="no 'choices'"): + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + with pytest.raises(APIError, match="no 'choices'"): + list(convert_to_streaming_response(response_object=resp)) + with pytest.raises(APIError, match="no 'choices'"): + async for _ in convert_to_streaming_response_async(response_object=resp): + pass diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index c59ec70b015..00b3a0633f2 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -41,6 +41,21 @@ from litellm.types.utils import ( ) +def test_translate_openai_response_to_anthropic_empty_choices() -> None: + response: Final = ModelResponse( + id="chatcmpl-empty", + model="gemini-3.5-flash", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10), + ) + + result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "end_turn" + assert result["usage"]["input_tokens"] == 10 + + def test_translate_chat_refusal_to_anthropic_response(): response = ModelResponse( id="chatcmpl-refusal", diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 038df3656fe..de6322f27ee 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -120,7 +120,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response + client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", From 261777d633c6c5466a0711a7a889ca2cfb0a08d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:40:29 -0700 Subject: [PATCH 116/157] fix(databricks): keep top-level reasoning_content from OpenAI-compatible gateway models The Databricks chat transformation only parsed reasoning out of FMAPI-style reasoning content blocks, so external models behind Databricks AI Gateway that return the OpenAI-style top-level reasoning_content string lost it, both in the final message and in every streamed delta. Fall back to the shared OpenAI reasoning helper when no reasoning block exists, and keep the delta's own reasoning_content when streaming. --- .../llms/databricks/chat/transformation.py | 33 ++++++-- litellm/types/llms/databricks.py | 9 ++ .../test_databricks_chat_transformation.py | 83 +++++++++++++++++++ 3 files changed, 118 insertions(+), 7 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e2e2ea5b553..54c7040daf9 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation strip_litellm_internal_message_fields, strip_name_from_message, ) @@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.databricks import ( AllDatabricksContentValues, DatabricksChoice, + DatabricksDelta, DatabricksFunction, + DatabricksMessage, DatabricksResponse, DatabricksTool, ) @@ -534,6 +537,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): thinking_blocks.append(thinking_block) return reasoning_content, thinking_blocks + @staticmethod + def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None: + return delta.get("reasoning_content") + + @staticmethod + def resolve_reasoning_and_content( + message: DatabricksMessage, block_reasoning_content: str | None + ) -> tuple[str | None, str | None]: + content_str: Final = DatabricksConfig.extract_content_str(message["content"]) + if block_reasoning_content is not None: + return block_reasoning_content, content_str + return _extract_reasoning_content({**message, "content": content_str}) + @staticmethod def extract_citations( content: AllDatabricksContentValues | None, @@ -577,14 +593,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): finish_reason = "stop" if translated_message is None: - ## get the content str - content_str = DatabricksConfig.extract_content_str(choice["message"]["content"]) - - ## get the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content")) + reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content( + choice["message"], block_reasoning_content + ) citations = DatabricksConfig.extract_citations(choice["message"].get("content")) @@ -738,12 +753,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): # extract the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content")) choice["delta"]["content"] = content_str - choice["delta"]["reasoning_content"] = reasoning_content + choice["delta"]["reasoning_content"] = ( + block_reasoning_content + if block_reasoning_content is not None + else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"]) + ) choice["delta"]["thinking_blocks"] = thinking_blocks translated_choices.append(choice) return ModelResponseStream( diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index e87a684aab8..a9c027bd2de 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -2,6 +2,7 @@ from typing import Any, Literal from pydantic import BaseModel from typing_extensions import ( + ReadOnly, Required, TypedDict, ) @@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False): role: Required[str] content: Required[AllDatabricksContentValues] tool_calls: list[DatabricksTool] | None + reasoning_content: ReadOnly[str | None] + reasoning: ReadOnly[str | None] + + +class DatabricksDelta(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[AllDatabricksContentValues | None] + reasoning_content: ReadOnly[str | None] class DatabricksChoice(TypedDict, total=False): diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 02655fb7f77..89d1aa52838 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -590,3 +590,86 @@ def test_chunk_parser_without_usage_still_parses_content(): assert result.id == "chatcmpl-test" assert result.model == "databricks-claude-sonnet-5" assert result.choices[0]["delta"]["content"] == "hi" + + +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_transform_choices_surfaces_top_level_reasoning_content(reasoning_key: str) -> None: + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": "391", + reasoning_key: "We need answer just number. 17*23=391.", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "We need answer just number. 17*23=391." + assert getattr(choices[0].message, "thinking_blocks", None) is None + + +def test_transform_choices_parses_think_tags_in_string_content(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": {"role": "assistant", "content": "17 times 23391"}, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "17 times 23" + + +def test_transform_choices_prefers_reasoning_blocks_over_top_level_field(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "from block"}]}, + {"type": "text", "text": "391"}, + ], + "reasoning_content": "from field", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.reasoning_content == "from block" + assert choices[0].message.content == "391" + + +def test_chunk_parser_surfaces_top_level_reasoning_delta(): + iterator = DatabricksChatResponseIterator(None, sync_stream=True) + chunk = { + "id": "1", + "object": "chat.completion.chunk", + "created": 0, + "model": "lit-qa-deepseek-v4-flash", + "choices": [ + { + "delta": {"role": "assistant", "content": "", "reasoning_content": "We need answer"}, + "index": 0, + "finish_reason": None, + } + ], + } + + parsed = iterator.chunk_parser(chunk) + + assert parsed.choices[0].delta.reasoning_content == "We need answer" + assert parsed.choices[0].delta.content == "" From 420282acb7704772ffb47c2ed4f84d455a2a385d Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 19:55:02 +0000 Subject: [PATCH 117/157] fix(registry): set max_output_tokens on vertex_ai/xai/grok-4.3 and grok-4.6 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 664c2737157..885021692c9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48204,6 +48204,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 200000, + "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, @@ -48222,6 +48223,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 524288, + "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 6e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 664c2737157..885021692c9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48204,6 +48204,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 200000, + "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, @@ -48222,6 +48223,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 524288, + "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 6e-06, From 17ca562b6a5d208d1b41a58933491936d4cefee2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:01:11 -0700 Subject: [PATCH 118/157] test(llm_translation): expect an empty choices list to convert instead of raising --- .../test_convert_dict_to_chat_completion.py | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index b6e30ddc711..09d1adba17c 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1623,15 +1623,11 @@ class TestMissingChoicesGuard: assert "no 'choices'" in exc_info.value.message - def test_convert_to_model_response_object_empty_choices_raises_api_error(self): - """Empty choices list raises APIError, same as missing/null choices. + def test_convert_to_model_response_object_empty_choices_returns_empty_list(self): + """An empty choices list is a real provider answer, so it converts to choices=[] instead of raising. - Provider-specific repair (e.g. github_copilot synthesizing choices for - Anthropic-native responses) happens before this guard, in the provider - config; the core utility keeps treating empty choices as an error. + See: https://github.com/BerriAI/litellm/issues/40276 """ - from litellm.exceptions import APIError - response_object = { "id": "msg_123", "model": "some-model", @@ -1639,13 +1635,14 @@ class TestMissingChoicesGuard: "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, } - with pytest.raises(APIError) as exc_info: - convert_to_model_response_object( - response_object=response_object, - model_response_object=ModelResponse(), - ) + result = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) - assert "no 'choices'" in exc_info.value.message + assert isinstance(result, ModelResponse) + assert result.choices == [] + assert result.usage.prompt_tokens == 10 def test_convert_to_model_response_object_null_choices_raises_api_error(self): """choices=None raises APIError.""" From f3a2844080f5b2dd003a084eece2861e78a1870d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:05:11 -0700 Subject: [PATCH 119/157] test(convert_dict_to_response): keep the regression test locals final and comment-free --- .../test_convert_dict_to_response.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index ea50bff462e..c3e99f01cec 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -108,7 +108,7 @@ def test_convert_empty_choices_response() -> None: convert_to_streaming_response, ) - resp = { + resp: Final = { "id": "x", "created": 1, "model": "gemini-3.5-flash", @@ -117,7 +117,7 @@ def test_convert_empty_choices_response() -> None: "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, "vertex_ai_safety_results": ["blocked"], } - result = convert_to_model_response_object( + result: Final = convert_to_model_response_object( response_object=resp, model_response_object=ModelResponse(), response_type="completion", @@ -125,8 +125,7 @@ def test_convert_empty_choices_response() -> None: assert result.choices == [] assert getattr(result, "vertex_ai_safety_results") == ["blocked"] - # Test sync streaming generator handles empty choices - sync_stream = list(convert_to_streaming_response(response_object=resp)) + sync_stream: Final = list(convert_to_streaming_response(response_object=resp)) assert len(sync_stream) == 1 assert sync_stream[0].choices == [] @@ -137,7 +136,7 @@ async def test_convert_empty_choices_response_async() -> None: convert_to_streaming_response_async, ) - resp = { + resp: Final = { "id": "x", "created": 1, "model": "gemini-3.5-flash", @@ -145,9 +144,7 @@ async def test_convert_empty_choices_response_async() -> None: "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, } - async_chunks = [] - async for chunk in convert_to_streaming_response_async(response_object=resp): - async_chunks.append(chunk) + async_chunks: Final = [chunk async for chunk in convert_to_streaming_response_async(response_object=resp)] assert len(async_chunks) == 1 assert async_chunks[0].choices == [] @@ -155,7 +152,7 @@ async def test_convert_empty_choices_response_async() -> None: def test_convert_missing_choices_raises_api_error() -> None: from litellm.exceptions import APIError - resp = { + resp: Final = { "id": "x", "created": 1, "model": "gemini-3.5-flash", From fe242f87028a9bd8ae2fcd4e579436f32614c116 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:16:57 -0700 Subject: [PATCH 120/157] test(databricks): pin the streamed reasoning delta shape and alias --- .../databricks/chat/test_databricks_chat_transformation.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 89d1aa52838..38d5844b9b9 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -653,7 +653,8 @@ def test_transform_choices_prefers_reasoning_blocks_over_top_level_field(): assert choices[0].message.content == "391" -def test_chunk_parser_surfaces_top_level_reasoning_delta(): +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> None: iterator = DatabricksChatResponseIterator(None, sync_stream=True) chunk = { "id": "1", @@ -662,7 +663,7 @@ def test_chunk_parser_surfaces_top_level_reasoning_delta(): "model": "lit-qa-deepseek-v4-flash", "choices": [ { - "delta": {"role": "assistant", "content": "", "reasoning_content": "We need answer"}, + "delta": {"role": "assistant", "content": None, reasoning_key: "We need answer"}, "index": 0, "finish_reason": None, } @@ -672,4 +673,4 @@ def test_chunk_parser_surfaces_top_level_reasoning_delta(): parsed = iterator.chunk_parser(chunk) assert parsed.choices[0].delta.reasoning_content == "We need answer" - assert parsed.choices[0].delta.content == "" + assert parsed.choices[0].delta.content is None From c4bd3763e87ededa07ac552bb5fab67691bd3d49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:17:19 -0700 Subject: [PATCH 121/157] test(proxy): type the background retrieval governance tests The two test methods, the policy_engine fixture, and the two inner stubs in TestBackgroundResponseRetrievalGovernance now carry full parameter and return annotations, closing the Greptile thread that 94f9230d13 left open. --- .../proxy/test_common_request_processing.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c451a3b4cb0..f119926e4f4 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Optional +from typing import AsyncGenerator, Callable, Final, Iterator, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -8303,7 +8303,7 @@ class TestBackgroundResponseRetrievalGovernance: GOVERNED_MODEL_ID = "deployment-governed" @pytest.fixture - def policy_engine(self): + def policy_engine(self) -> Iterator[None]: from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -8353,10 +8353,14 @@ class TestBackgroundResponseRetrievalGovernance: mock_request = MagicMock(spec=Request) mock_request.headers = {} - async def passthrough_add_litellm_data_to_request(data, **kwargs): + async def passthrough_add_litellm_data_to_request( + data: dict[str, object], **kwargs: object + ) -> dict[str, object]: return data - async def decrypting_pre_call_hook(user_api_key_dict, data, call_type): + async def decrypting_pre_call_hook( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: if data.get("response_id") == client_facing_response_id: data["response_id"] = encoded_response_id return data @@ -8383,8 +8387,8 @@ class TestBackgroundResponseRetrievalGovernance: @pytest.mark.asyncio async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( - self, policy_engine, monkeypatch - ): + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: data = await self._pre_call("aget_responses", monkeypatch) assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") @@ -8397,8 +8401,8 @@ class TestBackgroundResponseRetrievalGovernance: @pytest.mark.asyncio async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( - self, policy_engine, monkeypatch - ): + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: data = await self._pre_call("aresponses", monkeypatch) assert "_guardrail_pipelines" not in data["litellm_metadata"] From a9cce4f1ed15af0082a66233732d8e27f95b6bc8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:38:51 -0700 Subject: [PATCH 122/157] fix(guardrails): scan and rewrite Responses custom_tool_call output items Post-call guardrails on /v1/responses only treated function_call output items as tool calls, so a custom_tool_call item (Codex's exec shell tool on GPT-5.6 models) was never scanned or masked, non-streaming and streaming alike. Both item types now flow through the shared tool_call_dict_from_output_item helper, ended-stream delivery syncs the custom_tool_call_input delta/done events and the item's input field, and the completed-response scan key fingerprints both kinds of item. Non-streaming Responses tool-call MASK rewrites were also never written back to the output item even for function_call; they are now. --- .../transformation.py | 8 +- .../guardrail_translation/handler.py | 157 +++++++----- ...test_openai_responses_guardrail_handler.py | 229 +++++++++++++++++- 3 files changed, 322 insertions(+), 72 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 4fe069b0b7d..aa8337b6300 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): provider_specific_fields: Mapping[str, object] -def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: +def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw string payload in ``input`` rather than ``arguments``; both map to @@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Tool calls accumulate into the single trailing tool_calls choice # like the typed branches above; a choice per call would hide every # call after choices[0] from chat clients - accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index)) + accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index)) tool_call_index += 1 elif handle_raw_dict_callback is not None: choice, index = handle_raw_dict_callback(item=raw_item, index=index) @@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") in ("function_call", "custom_tool_call"): - converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) + converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) provider_specific_fields: Final = converted.get("provider_specific_fields") function_chunk: Final = ChatCompletionToolCallFunctionChunk( @@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): index=0, delta=Delta( tool_calls=( - _tool_call_dict_from_output_item( + tool_call_dict_from_output_item( output_item, parsed_chunk.get("output_index", 0) ), ) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 280a8670c36..e78e1f56915 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -37,7 +37,6 @@ from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -45,6 +44,7 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, + tool_call_dict_from_output_item, ) from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -84,7 +84,6 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import ( GenericResponseOutputItem, - OutputFunctionToolCall, OutputText, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -140,8 +139,18 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( ) -_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset( - {"response.function_call_arguments.delta", "response.function_call_arguments.done"} +_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) +_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call": "arguments", "custom_tool_call": "input"} +) +_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset( + {"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"} +) +_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"} +) +_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset( + _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS ) _OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( @@ -180,8 +189,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts -def _is_function_call_item(item: object) -> bool: - return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") +def _is_tool_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES + + +def _tool_call_output_item_mapping(item: object) -> Mapping[str, Any] | None: + if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES: + return None + if isinstance(item, Mapping): + return cast("Mapping[str, Any]", item) # cast-ok: output items are str-keyed JSON objects + return item.model_dump() if isinstance(item, BaseModel) else None + + +def _is_tool_call_output_item(item: object) -> bool: + return _tool_call_output_item_mapping(item) is not None def _last_message_role(messages: Sequence[object]) -> str | None: @@ -205,7 +226,7 @@ def _provenance_unit_bounds( start_indexes: Final = tuple( index for index in range(len(raw_input)) - if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") ) return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) @@ -603,7 +624,7 @@ class OpenAIResponsesHandler(BaseTranslation): - response.output is a list of output items - Each output item can be: * GenericResponseOutputItem with a content list of OutputText objects - * ResponseFunctionToolCall with tool call data + * ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data - Each OutputText object has a text field """ @@ -668,6 +689,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -676,6 +698,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -683,6 +706,15 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) + self._write_tool_call_rewrites_to_output( + tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=_tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) + else tool_calls_to_check + ), + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) @@ -933,11 +965,12 @@ class OpenAIResponsesHandler(BaseTranslation): guardrail_name: str, ) -> None: """Write ended-stream guardrail tool-call rewrites into the completed - envelope's ``function_call`` items and sync the earlier stream events, - keyed by ``call_id``. The guardrail sees the envelope's function calls - in output order, which is how a rewritten call finds its ``call_id``; - the stream events find their call through the ``call_id`` on - ``output_item`` events and the ``item_id`` on argument events, since an + envelope's ``function_call`` and ``custom_tool_call`` items and sync the + earlier stream events, keyed by ``call_id``. The guardrail sees the + envelope's tool calls in output order, which is how a rewritten call + finds its ``call_id``; the stream events find their call through the + ``call_id`` on ``output_item`` events and the ``item_id`` on argument + and custom-input events, since an event's ``output_index`` need not match the envelope's (the chat bridge numbers tool calls from 1 while the envelope lists them after the message). A rewrite whose calls do not line up with the envelope, or @@ -945,18 +978,16 @@ class OpenAIResponsesHandler(BaseTranslation): pipeline executor discards it and releases the original events.""" if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - function_call_items: Final = tuple( - output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call" - ) + tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item)) call_ids: Final = tuple( call_id - for output_item in function_call_items + for output_item in tool_call_items if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id ) stream_events: Final = responses_so_far[:-1] - call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events) + call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events) event_call_ids: Final = tuple( - self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events + self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events ) rewrites_by_call_id: Final = MappingProxyType( { @@ -966,11 +997,11 @@ class OpenAIResponsesHandler(BaseTranslation): } ) unresolved_argument_event: Final = any( - call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES + call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES for event, call_id in zip(stream_events, event_call_ids) ) if ( - len(call_ids) != len(function_call_items) + len(call_ids) != len(tool_call_items) or len(frozenset(call_ids)) != len(call_ids) or len(call_ids) != len(post_guardrail_tool_calls) or unresolved_argument_event @@ -981,10 +1012,10 @@ class OpenAIResponsesHandler(BaseTranslation): raise UndeliverableStreamRewrite(guardrail_name) for output_item, rewrite in ( (output_item, rewrites_by_call_id[call_id]) - for output_item, call_id in zip(function_call_items, call_ids) + for output_item, call_id in zip(tool_call_items, call_ids) if call_id in rewrites_by_call_id ): - self._write_function_call_item(output_item, rewrite.name, rewrite.arguments) + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) delta_replacements: Final = MappingProxyType( {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()} ) @@ -992,16 +1023,18 @@ class OpenAIResponsesHandler(BaseTranslation): if call_id not in rewrites_by_call_id: continue match stream_item_field(event, "type"): - case "response.function_call_arguments.delta": + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: self._write_event_field(event, "delta", next(delta_replacements[call_id])) - case "response.function_call_arguments.done": - self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments) + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: + self._write_event_field( + event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments + ) case "response.output_item.added": - self._write_function_call_item( + self._write_tool_call_item( stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None ) case "response.output_item.done": - self._write_function_call_item( + self._write_tool_call_item( stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, rewrites_by_call_id[call_id].arguments, @@ -1009,8 +1042,23 @@ class OpenAIResponsesHandler(BaseTranslation): case _: pass + def _write_tool_call_rewrites_to_output( + self, + tool_call_items: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + ) -> None: + if len(tool_call_items) != len(post_guardrail_tool_calls): + return + for output_item, after in ( + (output_item, after) + for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + ): + self._write_tool_call_item(output_item, after.name, after.arguments) + @staticmethod - def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: + def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: items: Final = tuple( stream_item_field(event, "item") for event in stream_events @@ -1020,32 +1068,35 @@ class OpenAIResponsesHandler(BaseTranslation): { item_id: call_id for item in items - if stream_item_field(item, "type") == "function_call" + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(item_id := stream_item_field(item, "id"), str) and isinstance(call_id := stream_item_field(item, "call_id"), str) } ) @staticmethod - def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: + def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: event_type: Final = stream_item_field(event, "type") - if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES: + if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES: item_id: Final = stream_item_field(event, "item_id") return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None if event_type not in _OUTPUT_ITEM_EVENT_TYPES: return None item: Final = stream_item_field(event, "item") call_id: Final = stream_item_field(item, "call_id") - return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None + return ( + call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None + ) @staticmethod - def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None: + def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None: if item is None: return if name is not None: OpenAIResponsesHandler._write_event_field(item, "name", name) - if arguments is not None: - OpenAIResponsesHandler._write_event_field(item, "arguments", arguments) + item_type: Final = stream_item_field(item, "type") + if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS: + OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload) def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ @@ -1073,7 +1124,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _completed_response_scan_key(response: object) -> StreamingScanKey: output_items: Final = stream_item_items(response, "output") message_items: Final = tuple( - item for item in output_items if stream_item_field(item, "type") != "function_call" + item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES ) return StreamingScanKey( texts=tuple( @@ -1085,7 +1136,7 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls=tuple( stream_item_fingerprint(item) for item in output_items - if stream_item_field(item, "type") == "function_call" + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES ), stream_ended=True, ) @@ -1196,34 +1247,10 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image/tool extraction logic. """ - # Check if this is a tool call (OutputFunctionToolCall) - if isinstance(output_item, OutputFunctionToolCall) or ( - isinstance(output_item, BaseModel) - and hasattr(output_item, "type") - and getattr(output_item, "type") == "function_call" - ): + tool_call_item: Final = _tool_call_output_item_mapping(output_item) + if tool_call_item is not None: if tool_calls_to_check is not None: - tool_call_dict = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=output_item, - index=output_idx, - ) - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - return - elif isinstance(output_item, dict) and output_item.get("type") == "function_call": - # Handle dict representation of tool call - if tool_calls_to_check is not None: - # Convert dict to ResponseFunctionToolCall for processing - try: - tool_call_obj: Final = ResponseFunctionToolCall(**output_item) - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=tool_call_obj, - index=output_idx, - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - except Exception: - pass + tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx)) return # Handle both GenericResponseOutputItem and dict diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 6ed4ec6618f..628ada05267 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -14,7 +14,12 @@ import pytest from fastapi import HTTPException -from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses import ( + ResponseCustomToolCall, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + ResponseFunctionToolCall, +) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -27,7 +32,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.responses.main import GenericResponseOutputItem, OutputText +from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -57,6 +62,37 @@ class MockGuardrail(CustomGuardrail): return inputs +class PersimmonMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": tool_call["function"]["arguments"].replace("persimmon", "[MASKED]"), + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + +CUSTOM_TOOL_CALL_ITEM = { + "type": "custom_tool_call", + "id": "ctc_1", + "call_id": "call_exec_1", + "name": "exec", + "input": "echo persimmon", + "status": "completed", +} + + class TestOpenAIResponsesHandlerDiscovery: """Test that the handler is properly discovered by the guardrail system""" @@ -628,6 +664,77 @@ class TestOpenAIResponsesHandlerToolCallExtraction: == '{"location":"Boston, MA","unit":"celsius"}' ) + @pytest.mark.parametrize( + "output_item", + [ + dict(CUSTOM_TOOL_CALL_ITEM), + CustomToolCallOutputItem(**CUSTOM_TOOL_CALL_ITEM), + ResponseCustomToolCall(**{key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "status"}), + ], + ids=["dict", "litellm_typed", "openai_typed"], + ) + def test_extract_custom_tool_call_input_as_arguments(self, output_item): + handler = OpenAIResponsesHandler() + texts_to_check: List[str] = [] + tool_calls_to_check: List[Any] = [] + + handler._extract_output_text_and_images( + output_item=output_item, + output_idx=2, + texts_to_check=texts_to_check, + images_to_check=[], + task_mappings=[], + tool_calls_to_check=tool_calls_to_check, + ) + + assert texts_to_check == [] + assert tool_calls_to_check == [ + { + "id": "call_exec_1", + "type": "function", + "function": {"name": "exec", "arguments": "echo persimmon"}, + "index": 2, + } + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("typed", [False, True], ids=["dict", "typed"]) + async def test_process_output_response_writes_tool_call_rewrites_back(self, typed): + handler = OpenAIResponsesHandler() + function_call = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_fn_1", + "name": "lookup_fruit", + "arguments": '{"fruit": "persimmon"}', + "status": "completed", + } + message = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "running persimmon", "annotations": []}], + } + payload = { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [message, function_call, dict(CUSTOM_TOOL_CALL_ITEM)], + } + response = ResponsesAPIResponse.model_validate(payload) if typed else payload + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + output = result.output if typed else result["output"] + function_item, custom_item = output[1], output[2] + assert (function_item.arguments if typed else function_item["arguments"]) == '{"fruit": "[MASKED]"}' + assert (custom_item.input if typed else custom_item["input"]) == "echo [MASKED]" + assert (custom_item.name if typed else custom_item["name"]) == "exec" + assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon" + @pytest.mark.asyncio async def test_process_output_response_with_tool_calls(self): """Test processing output response containing function tool calls""" @@ -1315,6 +1422,109 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' assert completed_event.response.output[0].name == "lookup_fruit" + @staticmethod + def _ended_custom_tool_call_stream_events() -> List[dict]: + def item(input_text: str, status: str) -> dict: + return {**CUSTOM_TOOL_CALL_ITEM, "input": input_text, "status": status} + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"}, + {"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"}, + {"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "created_at": 1, + "model": "gpt-5.6", + "output": [item("echo persimmon", "completed")], + "status": "completed", + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_custom_tool_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["input"] == "" + assert events[1]["delta"] == "echo [MASKED]" + assert events[2]["delta"] == "" + assert events[3]["input"] == "echo [MASKED]" + assert events[4]["item"]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["name"] == "exec" + assert "arguments" not in events[5]["response"]["output"][0] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self): + from litellm.types.llms.openai import ( + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[Any] = [ + model.model_validate({**event, "sequence_number": sequence_number}) + for sequence_number, (model, event) in enumerate( + zip( + ( + OutputItemAddedEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_custom_tool_call_stream_events(), + ) + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event.response.output[0], CustomToolCallOutputItem) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == "echo [MASKED]" + assert typed_events[2].delta == "" + assert typed_events[3].input == "echo [MASKED]" + assert typed_events[4].item.input == "echo [MASKED]" + assert completed_event.response.output[0].input == "echo [MASKED]" + assert completed_event.response.output[0].name == "exec" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @staticmethod def _bridged_function_call_stream_events() -> List[dict]: reasoning = {"type": "reasoning", "id": "rs_1", "summary": []} @@ -2747,8 +2957,21 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + def test_completed_event_with_a_custom_tool_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + ended_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, dict(CUSTOM_TOOL_CALL_ITEM)])] + ) + rewritten_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, {**CUSTOM_TOOL_CALL_ITEM, "input": "echo kumquat"}])] + ) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "echo persimmon" in ended_key.tool_calls[0] + assert rewritten_key != ended_key + def test_completed_event_reads_every_output_text_part(self): - from litellm.types.responses.main import GenericResponseOutputItem, OutputText + from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText item = GenericResponseOutputItem( type="message", From c5e93aff132685582679ed1ce7097e7e855630a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:45:26 -0700 Subject: [PATCH 123/157] fix(proxy): warn at submit when a body-selected post_call policy is deferred --- litellm/proxy/utils.py | 17 ++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 31 +++++++++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1604dea1003..1f8a4d5ba33 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -622,6 +622,15 @@ def _defer_post_call_pipelines( response.id, ", ".join(tag_matched), ) + body_selected: Final = _body_selected_deferrals(data, deferred) + if body_selected: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: %s", + response.id, + ", ".join(body_selected), + ) _withdraw_deferred_claims(data, deferred) @@ -638,6 +647,14 @@ def _tag_matched_deferrals( ) +def _body_selected_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset() + return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed) + + def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 3bbdbbe2ccd..cc40706c67c 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -941,6 +941,7 @@ def _post_call_pipeline_data( "metadata": { "_guardrail_pipelines": [("response-governance", pipeline)], "_pipeline_managed_guardrails": {guardrail}, + "policy_sources": {"response-governance": "model:m"}, }, **extra, } @@ -1468,7 +1469,7 @@ def _output_passing_callbacks() -> list[CustomGuardrail]: def _claimed_post_call_pipeline_data( - *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str = "model:m" + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m" ): from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -1491,7 +1492,7 @@ def _claimed_post_call_pipeline_data( "_pipeline_managed_guardrails": {"gr-post"}, "applied_policies": list(policy_names), "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], - "policy_sources": {policy_name: policy_source for policy_name in policy_names}, + "policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None}, }, } @@ -1551,6 +1552,32 @@ async def test_pending_background_response_warns_when_the_deferred_policy_was_ma ] +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("body-governance", policy_source=None) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert _warnings(caplog) == [ + "Policy engine: background response resp_bg matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: body-governance" + ] + + @pytest.mark.asyncio async def test_pending_background_response_matched_through_its_model_does_not_warn( proxy_logging: ProxyLogging, From 5bdf45726d0a27372c1cc34ee6baa85ef5439650 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:52:08 -0700 Subject: [PATCH 124/157] fix(bedrock): honor the json_mode kwarg on invoke Nova structured output and keep gpt-4o-mini off chat web search --- .../bedrock/chat/converse_transformation.py | 10 +++-- ...odel_prices_and_context_window_backup.json | 6 +-- model_prices_and_context_window.json | 6 +-- .../chat/test_converse_transformation.py | 38 +++++++++++++++++++ whitelisted_bedrock_models.txt | 1 + 5 files changed, 50 insertions(+), 11 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index fa24f8be893..6173f08fc87 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1805,6 +1805,7 @@ class AmazonConverseConfig(BaseConfig): data=request_data, messages=messages, encoding=encoding, + json_mode=json_mode, ) def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str: @@ -2237,6 +2238,7 @@ class AmazonConverseConfig(BaseConfig): data: dict | str, messages: list, encoding, + json_mode: bool | None = None, ) -> ModelResponse: ## LOGGING if logging_obj is not None: @@ -2247,7 +2249,9 @@ class AmazonConverseConfig(BaseConfig): additional_args={"complete_input_dict": data}, ) - json_mode: Final[bool | None] = optional_params.get("json_mode", None) + resolved_json_mode: Final[bool | None] = ( + json_mode if json_mode is not None else optional_params.get("json_mode", None) + ) ## RESPONSE OBJECT try: completion_response: Final = ConverseResponseBlock(**response.json()) @@ -2339,7 +2343,7 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks) chat_completion_message["content"] = content_str filtered_tools: Final = self._filter_json_mode_tools( - json_mode=json_mode, + json_mode=resolved_json_mode, tools=tools, chat_completion_message=chat_completion_message, ) @@ -2363,7 +2367,7 @@ class AmazonConverseConfig(BaseConfig): # When json_mode filtered out all synthetic tool calls the response # is plain content, not a pending tool invocation. Fix finish_reason # so callers (e.g. OpenAI SDK) don't misinterpret it. - if json_mode and not filtered_tools and tools: + if resolved_json_mode and not filtered_tools and tools: initial_finish_reason = "stop" ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 885021692c9..a2e7b692649 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29407,8 +29407,7 @@ "search_context_size_high": 0.025, "search_context_size_low": 0.025, "search_context_size_medium": 0.025 - }, - "supports_web_search": true + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29436,8 +29435,7 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_vision": true }, "gpt-4o-mini-audio-preview": { "deprecation_date": "2026-05-07", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 885021692c9..a2e7b692649 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29407,8 +29407,7 @@ "search_context_size_high": 0.025, "search_context_size_low": 0.025, "search_context_size_medium": 0.025 - }, - "supports_web_search": true + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29436,8 +29435,7 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true + "supports_vision": true }, "gpt-4o-mini-audio-preview": { "deprecation_date": "2026-05-07", diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index f0e361ceb88..16824d19f23 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6727,3 +6727,41 @@ def test_forced_tool_choice_forwarded_on_converse_models_that_support_it( ) assert result == {"any": {}} + + +def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it(): + response_json = { + "metrics": {"latencyMs": 900}, + "output": { + "message": { + "content": [ + { + "toolUse": { + "input": {"city": "Paris", "population": 2100000}, + "name": "json_tool_call", + "toolUseId": "tooluse_invoke_nova_json", + } + } + ], + "role": "assistant", + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 40, "outputTokens": 20, "totalTokens": 60}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock.test")) + logging_obj = MagicMock() + result = AmazonConverseConfig().transform_response( + model="bedrock/invoke/us.amazon.nova-micro-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": {}}}]}, + litellm_params={}, + encoding=None, + json_mode=True, + ) + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt index 578edddec8d..6124cb41044 100644 --- a/whitelisted_bedrock_models.txt +++ b/whitelisted_bedrock_models.txt @@ -6,6 +6,7 @@ ai21.jamba-instruct-v1:0 twelvelabs.pegasus-1-2-v1:0 us.twelvelabs.pegasus-1-2-v1:0 eu.twelvelabs.pegasus-1-2-v1:0 +global.twelvelabs.pegasus-1-2-v1:0 amazon.titan-text-express-v1 amazon.titan-text-lite-v1 amazon.titan-text-premier-v1:0 From b89c32407f81141e6c877b534862ad5b3e7956e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:01:36 -0700 Subject: [PATCH 125/157] test(cost): assert gpt-4o-mini keeps chat web search off while the flat search fee still matches its alias --- .../llm_cost_calc/test_tool_call_cost_tracking.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index f319413045f..bbb7b5f9c35 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -879,8 +879,8 @@ def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( alias_info = litellm.get_model_info("gpt-4o-mini") snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18") - assert snapshot_info["supports_web_search"] is True - assert alias_info["supports_web_search"] is True + assert not snapshot_info["supports_web_search"] + assert not alias_info["supports_web_search"] snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( web_search_options=web_search_options, model_info=snapshot_info From 29af8b734913f5e43e4149bcdcc36a7c9261cc37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:04:54 -0700 Subject: [PATCH 126/157] fix(streaming_handler): replay a cached completion with no choices as an empty stream A stream cache hit on an entry stored with choices == [] indexed choices[0] in the cached_response branch and failed with IndexError, so the streaming converters' empty chunk had no working consumer. The branch now treats a chunk without choices as empty and lets the wrapper close the stream with its usual finish_reason stop chunk --- .../litellm_core_utils/streaming_handler.py | 11 ++--- .../test_streaming_handler.py | 43 ++++++++++++++++++- 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ae17bac6ff..db23929e0c3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1473,17 +1473,14 @@ class CustomStreamWrapper: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": cached_chunk: Final = cast(ModelResponseStream, chunk) - chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason + cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None + chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None response_obj = { - "text": cached_chunk.choices[0].delta.content, + "text": cached_choice.delta.content if cached_choice is not None else None, "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, "original_chunk": cached_chunk, - "tool_calls": ( - cached_chunk.choices[0].delta.tool_calls - if hasattr(cached_chunk.choices[0].delta, "tool_calls") - else None - ), + "tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None), } completion_obj["content"] = response_obj["text"] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 0aa73833677..ced7d6d677d 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -6,7 +6,7 @@ import pytest import asyncio import traceback -from typing import Optional +from typing import Final, Optional import litellm from litellm import verbose_logger @@ -2633,6 +2633,47 @@ def test_dispatch_cached_response_extracts_delta( assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1" +def test_dispatch_cached_response_without_choices_is_an_empty_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A cached completion with no choices replays as an empty, unfinished chunk + instead of raising IndexError on choices[0].""" + initialized_custom_stream_wrapper.custom_llm_provider = "cached_response" + chunk: Final = ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + result, model_response, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] is None + assert initialized_custom_stream_wrapper.received_finish_reason is None + assert model_response.id == "chatcmpl-cache-empty" + + +@pytest.mark.asyncio +async def test_cached_response_without_choices_streams_a_single_stop_chunk( + logging_obj: Logging, +): + """A stream cache hit on a completion stored with choices == [] ends with one + finish_reason=stop chunk, the same shape the live empty stream produced.""" + + async def cached_chunks(): + yield ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + wrapper: Final = CustomStreamWrapper( + completion_stream=cached_chunks(), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="cached_response", + ) + + chunks: Final = tuple([chunk async for chunk in wrapper]) + + assert tuple(choice.finish_reason for chunk in chunks for choice in chunk.choices) == ("stop",) + assert all(choice.delta.content in (None, "") for chunk in chunks for choice in chunk.choices) + + def test_dispatch_vertex_ai_legacy_text_and_finish_reason( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 2c836f473c8ef6574ae7a04ffa8434f7715e7e2f Mon Sep 17 00:00:00 2001 From: "cursor[bot]" <206951365+cursor[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:14:24 -0700 Subject: [PATCH 127/157] test(ui): derive reasoning-effort assertion from the anthropic preset (#40456) Preset #40341 pointed the Anthropic family REASONING tier at claude-fable-5-1, but this test still hardcoded claude-opus-5, so the payload it saw no longer matched. Rebase the assertion on ANTHROPIC_PRESET.complexity_router_config.tier_model_configs so a preset refresh flows through instead of redding the suite on staging. Co-authored-by: Cursor Agent Co-authored-by: Krrish Dholakia --- .../src/components/add_model/add_auto_router_tab.test.tsx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 014854ac712..f6e619c5e84 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1046,9 +1046,7 @@ describe("AddAutoRouterTab", () => { await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ complexity_router_config: { - tier_model_configs: { - REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }], - }, + tier_model_configs: ANTHROPIC_PRESET.complexity_router_config.tier_model_configs, }, }); }); From 354365eecaf6c70aba2a649fa324bf0369438f85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:15:01 -0700 Subject: [PATCH 128/157] test(proxy): give the pipeline-managed native hook audit test its chat completions route --- tests/test_litellm/proxy/test_proxy_logging_hook_detection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 42888ca7d54..28ff4571b44 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -688,7 +688,7 @@ async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monk "messages": [{"role": "user", "content": "hi"}], "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, }, - captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), captured_logging_obj=_streaming_logging_obj(), assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), cache_hit=False, From 0a053d2c8146e4a4739f92d48f1e2bdec75202e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:21:40 -0700 Subject: [PATCH 129/157] test(guardrails): cover streamed tool-call name rewrites on chat and Messages Skipping the name write-back in either handler left every test green; a guardrail that renames a tool call now has a regression test on both the chat chunk path and the Anthropic SSE path --- .../test_anthropic_guardrail_handler.py | 23 +++++++++++++++++++ .../test_openai_guardrail_handler.py | 23 +++++++++++++++++++ 2 files changed, 46 insertions(+) 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 64c243343d5..e091355b69d 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 @@ -367,6 +367,29 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"stop_reason": "tool_use"' in raw assert "persim" not in raw + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.name = "lookup_fruit_reviewed" + return inputs + + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw + assert '"name": "lookup_fruit"' not in raw + assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + @pytest.mark.asyncio async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index c7011a6f00e..5a29a96829f 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1171,6 +1171,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[3].choices[0].delta.tool_calls is None assert chunks[3].choices[0].finish_reason == "tool_calls" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call["function"]["name"] = "lookup_fruit_reviewed" + return inputs + + handler = OpenAIChatCompletionsHandler() + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]] + assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None] + assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"} + assert fragments[0].id == "call_1" + @pytest.mark.asyncio async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): handler = OpenAIChatCompletionsHandler() From cce1d2087b9ae997ec406d6d38857642f1461844 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:22:22 -0700 Subject: [PATCH 130/157] fix(responses): echo a named tool_choice in the Responses API shape on the chat-completions bridge A streamed /v1/responses request with tool_choice {"type": "function", "name": ...} that reaches a chat-completions-only deployment failed with HTTP 500 before the first byte: the synthetic response.created and response.in_progress events copied the chat-shaped tool_choice into ResponsesAPIResponse, whose ToolChoice type expects the flat Responses API shape. The non-streamed path echoed "auto" regardless of the request. Both paths now normalize the request's tool_choice through the existing chat transform and map it back to the Responses API vocabulary, validated by a TypeAdapter(ToolChoice), so a named function is echoed as {"type": "function", "name": ...} and a missing tool_choice is echoed as "auto". Fixes #33689 --- .../streaming_iterator.py | 11 ++-- .../transformation.py | 18 ++++++- .../test_litellm_completion_responses.py | 52 ++++++++++++++++++ .../test_streaming_iterator_transformation.py | 53 +++++++++++++++++++ 4 files changed, 126 insertions(+), 8 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index d7f8cd8f8bd..660dd8f0c92 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): response_created_event_data["temperature"] = self.responses_api_request["temperature"] if "text" in self.responses_api_request: response_created_event_data["text"] = self.responses_api_request["text"] - if "tool_choice" in self.responses_api_request: - # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format - response_created_event_data["tool_choice"] = ( - LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"]) - or "auto" + response_created_event_data["tool_choice"] = ( + LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + self.responses_api_request.get("tool_choice") ) - else: - response_created_event_data["tool_choice"] = "auto" + ) if "tools" in self.responses_api_request: response_created_event_data["tools"] = self.responses_api_request["tools"] else: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..878f493b58e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,6 +27,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -68,6 +69,7 @@ from litellm.types.llms.openai import ( ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStatus, + ToolChoice, ValidChatCompletionMessageContentTypes, ValidChatCompletionMessageContentTypesLiteral, ) @@ -126,6 +128,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) +_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @runtime_checkable @@ -267,6 +270,17 @@ class LiteLLMCompletionResponsesConfig: # Return as-is for unknown formats return tool_choice + @staticmethod + def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: + normalized: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) + match normalized: + case None: + return "auto" + case {"type": "function", "function": {"name": str(function_name)}}: + return ToolChoiceFunctionParam(type="function", name=function_name) + case _: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) + @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: """ @@ -2263,7 +2277,9 @@ class LiteLLMCompletionResponsesConfig: ), parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False), temperature=getattr(chat_completion_response, "temperature", 0), - tool_choice=getattr(chat_completion_response, "tool_choice", "auto"), + tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + responses_api_request.get("tool_choice") + ), tools=getattr(chat_completion_response, "tools", []), top_p=getattr(chat_completion_response, "top_p", None), max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2068f10ea2d..2aab660b5b7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1421,6 +1421,58 @@ class TestToolChoiceTransformation: ) assert result == "required" + @pytest.mark.parametrize( + "request_tool_choice,expected", + [ + ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), + ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "function", "name": "ApplyPatch"}), + ({"type": "tool"}, "required"), + ({"type": "auto"}, "auto"), + ("required", "required"), + ("none", "none"), + (None, "auto"), + ], + ) + def test_transform_tool_choice_for_responses_api_response(self, request_tool_choice, expected): + result = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + request_tool_choice + ) + assert result == expected + + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self): + chat_completion_response = ModelResponse( + id="chatcmpl-named-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_pwd", + type="function", + function=Function(name="run_command", arguments='{"command":"pwd"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 719d51c11e3..d4b565f82a1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -628,3 +628,56 @@ def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id(): assert item_dones[0].item.call_id == "toolu_01AbCdEf" for evt in deltas + dones: assert evt.item_id == added[0].item.id + + +def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_pwd", + "type": "function", + "function": {"name": "run_command", "arguments": '{"command":"pwd"}'}, + "index": 0, + } + ], + ), + finish_reason=finish_reason, + ) + ], + ) + + +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): + iterator = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": {"type": "function", "name": "run_command"}, + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + events = list(iterator) + + response_events = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + assert [event.type for event in response_events] == [ + "response.created", + "response.in_progress", + "response.completed", + ] + assert response_events[0].response.tool_choice == {"type": "function", "name": "run_command"} + assert response_events[1].response.tool_choice == {"type": "function", "name": "run_command"} + assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) From 5e7cb898065158626fec3a4afb90cb23a411095c Mon Sep 17 00:00:00 2001 From: Praveena Date: Tue, 4 Aug 2026 16:33:28 +0530 Subject: [PATCH 131/157] fix(azure): respect DEFAULT_MAX_RETRIES in initialize_azure_sdk_client When litellm_params does not include max_retries (the common case for deployments configured without explicit retry settings), initialize_azure_sdk_client() previously passed None through the 'if max_retries is not None' guard, resulting in AsyncAzureOpenAI being created without a max_retries argument. The OpenAI SDK then uses its own hardcoded default of 2, ignoring the DEFAULT_MAX_RETRIES env var. Fix: fall back to litellm.constants.DEFAULT_MAX_RETRIES when litellm_params has no max_retries. This ensures the SDK client respects the configured retry count. Steps to reproduce: 1. Set env var DEFAULT_MAX_RETRIES=0 2. Configure a deployment without explicit max_retries in litellm_params 3. Make a request that triggers a timeout 4. Observe: SDK retries (x-stainless-retry-count=1) despite env var=0 Related: https://github.com/BerriAI/litellm/issues/5124 --- litellm/llms/azure/common_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..f20cc37572b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -582,7 +582,11 @@ class BaseAzureLLM(BaseOpenAILLM): if scope is None: scope = "https://cognitiveservices.azure.com/.default" - max_retries: Final = litellm_params.get("max_retries") + max_retries = litellm_params.get("max_retries") + if max_retries is None: + from litellm.constants import DEFAULT_MAX_RETRIES + + max_retries = DEFAULT_MAX_RETRIES timeout: Final = litellm_params.get("timeout") if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret: verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") From 830ee2d39d84eba4d2c8a27851e363473af038d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:29:06 -0700 Subject: [PATCH 132/157] fix(azure): default max_retries to DEFAULT_MAX_RETRIES with regression tests initialize_azure_sdk_client now falls back to litellm.constants.DEFAULT_MAX_RETRIES when litellm_params carries no max_retries, so off-router Azure clients (files, batches, fine-tuning, assistants, audio) honor the env var like OpenAI clients do. Router paths already default max_retries to 0 and are unchanged. Regression tests cover the default, explicit 0/5/None values, and the env var reaching the SDK client in a fresh interpreter. --- litellm/llms/azure/common_utils.py | 11 ++-- .../llms/azure/test_azure_common_utils.py | 52 +++++++++++++++++++ 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index f20cc37572b..e6b3eb1f2bb 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import BaseOpenAILLM from litellm.secret_managers.get_azure_ad_token_provider import ( @@ -582,11 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM): if scope is None: scope = "https://cognitiveservices.azure.com/.default" - max_retries = litellm_params.get("max_retries") - if max_retries is None: - from litellm.constants import DEFAULT_MAX_RETRIES - - max_retries = DEFAULT_MAX_RETRIES + configured_max_retries: Final = litellm_params.get("max_retries") + max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries timeout: Final = litellm_params.get("timeout") if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret: verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") @@ -646,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM): else: azure_client_params["http_client"] = self._get_sync_http_client() - if max_retries is not None: - azure_client_params["max_retries"] = max_retries + azure_client_params["max_retries"] = max_retries if timeout is not None: azure_client_params["timeout"] = timeout diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..c959c201ccb 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -385,6 +385,58 @@ def test_select_azure_base_url_called(setup_mocks): setup_mocks["select_url"].assert_called_once() +def test_initialize_defaults_max_retries_to_litellm_default(setup_mocks): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == litellm.constants.DEFAULT_MAX_RETRIES + + +@pytest.mark.parametrize( + "configured, expected", + [(0, 0), (5, 5), (None, litellm.constants.DEFAULT_MAX_RETRIES)], +) +def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expected): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={"max_retries": configured}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == expected + + +def test_default_max_retries_env_var_reaches_azure_sdk_client(): + import subprocess + import sys + + code = ( + "from litellm.llms.azure.common_utils import BaseAzureLLM\n" + "client = BaseAzureLLM().get_azure_openai_client(" + "api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01'," + " client=None, _is_async=True, litellm_params={}, model='gpt-4')\n" + "print(client.max_retries)" + ) + completed = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DEFAULT_MAX_RETRIES": "0"}, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout.strip() == "0" + + @pytest.mark.parametrize( "call_type", [ From 2488f84b0293956b3736fcd9cf639491ab68e65e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:53:35 -0700 Subject: [PATCH 133/157] fix(proxy): keep a body litellm_session_id in SpendLogs under missing_session_id omit (#40379) * fix(proxy): keep a body litellm_session_id in SpendLogs under missing_session_id omit Under general_settings.missing_session_id: omit, apply_missing_session_id_policy now mirrors a client-supplied top-level litellm_session_id into metadata.session_id when the client did not set one there, so SpendLogs.session_id and Langfuse agree with the session callbacks already report through StandardLoggingPayload.session_id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep client metadata.session_id ahead of body litellm_session_id on litellm_metadata routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop docstrings from the missing_session_id omit regression tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 10 ++ .../proxy/test_litellm_pre_call_utils.py | 101 ++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3250ae5cca9..2925226f235 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -773,6 +773,16 @@ def apply_missing_session_id_policy( return if policy == "omit": metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + requester_metadata: Final = data.get("metadata") + requester_session_id: Final = ( + requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None + ) + if ( + (body_session_id := data.get("litellm_session_id")) + and not metadata.get("session_id") + and not requester_session_id + ): + metadata["session_id"] = body_session_id return if data.get("litellm_session_id") or metadata.get("session_id"): return 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 892fa484ab4..7564fa91c43 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7678,6 +7678,107 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id(): assert _spend_log_session_id(updated) == "client-session-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_body", + [ + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1"}, + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1", "metadata": {"trace_id": "trace-1"}}, + ], +) +async def test_missing_session_id_omit_keeps_body_litellm_session_id( + monkeypatch: pytest.MonkeyPatch, client_body: dict[str, object] +): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + updated = await add_litellm_data_to_request( + data=client_body, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + callback_session_id = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=SimpleNamespace(litellm_session_id=""), + litellm_params=get_litellm_params(litellm_session_id="cust-sess-1", metadata=updated["metadata"]), + ) + assert callback_session_id == "cust-sess-1" + assert updated["metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated) == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_body_litellm_session_id_does_not_override_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "messages": [], + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated) == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_metadata_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "input": "hi", + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_body_litellm_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "input": "hi", "litellm_session_id": "cust-sess-1"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_ignores_empty_body_litellm_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "litellm_session_id": ""}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert _spend_log_session_id(updated) is None + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" From cbeac276561c82e71f56de6d76e4d132170a2cc6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:12:13 -0700 Subject: [PATCH 134/157] fix(convert_dict_to_response): name the non-list choices type in the converter error When a provider returns choices as null, an object, a string or a number, the converter said the response had no 'choices' even though the key was present in the raw keys it listed. A shared message now keeps the old wording for a missing key and names the offending type otherwise. The cached-stream regression test also pins the chunk count so a leaked extra chunk fails it. --- .../convert_dict_to_response.py | 25 +++++++++++-------- .../test_convert_dict_to_response.py | 11 ++++---- .../test_streaming_handler.py | 1 + 3 files changed, 21 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 29b75812bfb..87524d86c61 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,7 @@ import json import re import time import traceback -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast import litellm @@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: del choice.enhancements +def _invalid_choices_message(response_object: Mapping[str, object]) -> str: + raw_keys: Final = list(response_object.keys()) + if "choices" not in response_object: + return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}" + return ( + f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). " + f"Raw keys: {raw_keys}" + ) + + async def convert_to_streaming_response_async( response_object: dict | None = None, ): @@ -184,9 +194,7 @@ async def convert_to_streaming_response_async( raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -292,9 +300,7 @@ def convert_to_streaming_response( raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -628,10 +634,7 @@ def convert_to_model_response_object( raise APIError( status_code=500, - message=( - "LiteLLM: provider returned a response with no 'choices'. " - f"Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index c3e99f01cec..8e46ae21de6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -167,9 +167,9 @@ def test_convert_missing_choices_raises_api_error() -> None: assert "no 'choices'" in str(exc_info.value) -@pytest.mark.parametrize("choices", [{}, "", None, 0]) +@pytest.mark.parametrize(("choices", "type_name"), [({}, "dict"), ("", "str"), (None, "NoneType"), (0, "int")]) @pytest.mark.asyncio -async def test_convert_non_list_choices_raises_api_error(choices: object) -> None: +async def test_convert_non_list_choices_raises_api_error(choices: object, type_name: str) -> None: from litellm.exceptions import APIError from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_streaming_response, @@ -183,14 +183,15 @@ async def test_convert_non_list_choices_raises_api_error(choices: object) -> Non "object": "chat.completion", "choices": choices, } - with pytest.raises(APIError, match="no 'choices'"): + expected: Final = f"'choices' that is not a list \\({type_name}\\)" + with pytest.raises(APIError, match=expected): convert_to_model_response_object( response_object=resp, model_response_object=ModelResponse(), response_type="completion", ) - with pytest.raises(APIError, match="no 'choices'"): + with pytest.raises(APIError, match=expected): list(convert_to_streaming_response(response_object=resp)) - with pytest.raises(APIError, match="no 'choices'"): + with pytest.raises(APIError, match=expected): async for _ in convert_to_streaming_response_async(response_object=resp): pass diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index ced7d6d677d..37e2031fdf4 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2670,6 +2670,7 @@ async def test_cached_response_without_choices_streams_a_single_stop_chunk( chunks: Final = tuple([chunk async for chunk in wrapper]) + assert len(chunks) == 1 assert tuple(choice.finish_reason for chunk in chunks for choice in chunk.choices) == ("stop",) assert all(choice.delta.content in (None, "") for chunk in chunks for choice in chunk.choices) From 264fc82dc1bb66e840b625830c9d8be03d655a0b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:14:10 -0700 Subject: [PATCH 135/157] fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages (#37228) * fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(playground): honor the Stream responses toggle for /v1/messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rate_limiter): drop explanatory docstrings from v3 dict response tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 18 ++--- .../hooks/parallel_request_limiter_v3.py | 33 +++------- .../add_retry_fallback_headers.py | 6 ++ .../hooks/test_dynamic_rate_limiter_v3.py | 57 ++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 65 ++++++++++++++++++ .../components/chat_ui/ChatUI.test.tsx | 53 +++++++++++++++ .../playground/components/chat_ui/ChatUI.tsx | 6 +- .../llm_calls/anthropic_messages.test.tsx | 66 ++++++++++++++++++- .../llm_calls/anthropic_messages.tsx | 33 +++++++--- 9 files changed, 291 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index de8834449de..a074f02f4e8 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Add additional priority-specific headers - if isinstance(response, ModelResponse): + if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) - - # Get existing additional headers - additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} - - # Add priority information + additional_headers: Final = ensure_response_additional_headers(response) additional_headers["x-litellm-priority"] = priority or "default" additional_headers["x-litellm-rate-limiter-version"] = "v3" - # Update response - if not hasattr(response, "_hidden_params"): - response._hidden_params = {} - response._hidden_params["additional_headers"] = additional_headers - return response except Exception as e: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..a72ae3bb1ea 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Post-call hook to update rate limit headers in the response. """ try: - from pydantic import BaseModel - stash: Final = get_request_stash() litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None - if litellm_proxy_rate_limit_response is not None: - # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") - else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() - - _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, + if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers.update( + self._merge_ratelimit_statuses_into_additional_headers( + additional_headers={}, statuses=litellm_proxy_rate_limit_response["statuses"], ) - - setattr( - response, - "_hidden_params", - {**_hidden_params, "additional_headers": _additional_headers}, - ) + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 3ec92ad226a..3251ea457cf 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None: return response +def response_has_hidden_params(response: object) -> bool: + if isinstance(response, dict): + return "_hidden_params" in response + return hasattr(response, "_hidden_params") + + def ensure_response_additional_headers(response: object) -> dict[str, object]: hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict)) _write_hidden_params(response, hidden_params) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0ff8b67b1a7..0cd6b4ede9c 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1861,3 +1861,60 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): ) assert capacity_blocked.value.status_code == 429 assert "Model capacity reached" in capacity_blocked.value.detail["error"] + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_priority_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RateLimitResponse, + RateLimitStatus, + get_or_create_request_stash, + ) + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=75, + limit_remaining=74, + rate_limit_type="requests", + descriptor_key="priority_model", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75 + assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74 + assert additional_headers["x-litellm-priority"] == "premium" + assert additional_headers["x-litellm-rate-limiter-version"] == "v3" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 4003286d887..6d382370f5f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6171,3 +6171,68 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses(): data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5)) ) assert get_request_stash().batch_enqueued_reservation == reservation + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100 + assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99 + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index f07b66efdf8..bf6b092a2b0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages"; vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), @@ -14,6 +15,10 @@ vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined), })); +vi.mock("../../llm_calls/anthropic_messages", () => ({ + makeAnthropicMessagesRequest: vi.fn().mockResolvedValue(undefined), +})); + vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({}), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), @@ -32,6 +37,8 @@ beforeEach(() => { const CHAT_REQUEST_ARG_COUNT = 26; const STREAMING_ENABLED_ARG_INDEX = 25; +const MESSAGES_REQUEST_ARG_COUNT = 19; +const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18; async function openComboboxByPlaceholder(placeholder: string) { const user = userEvent.setup(); @@ -378,6 +385,52 @@ describe("ChatUI", () => { expect(requestArgs[STREAMING_ENABLED_ARG_INDEX]).toBe(false); }); + it("should send the /v1/messages request non-streaming after Stream responses is unchecked", async () => { + const user = userEvent.setup(); + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/messages"); + await selectComboboxOption("Select a Model", "Model 1"); + + await user.click(await screen.findByTestId("model-settings-button")); + + const streamingCheckbox = await screen.findByRole("checkbox", { name: /Stream responses/i }); + expect(streamingCheckbox).toBeChecked(); + await user.click(streamingCheckbox); + + await waitFor(() => { + expect(screen.getByRole("checkbox", { name: /Stream responses/i })).not.toBeChecked(); + }); + + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + }); + await act(async () => { + fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" }); + }); + + await waitFor(() => { + expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1); + }); + + const requestArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0]; + expect(requestArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT); + expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false); + }); + it("should force streaming in simplified mode even when the playground setting is off", async () => { sessionStorage.setItem("streamingEnabled", "false"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 35378d3d4e7..eca9e2323ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -1025,6 +1025,7 @@ const ChatUI: React.FC = ({ mcpServers, mcpServerToolRestrictions, mcpToolsets, + streamingEnabled, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( @@ -1174,7 +1175,10 @@ const ChatUI: React.FC = ({ return !model.mode || model.mode === "chat"; }; - const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + const supportsStreamingToggle = + endpointType === EndpointType.CHAT || + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES; const modelsForEndpoint = useMemo( () => filterModelsForEndpoint(modelInfo, endpointType as EndpointType), [modelInfo, endpointType], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx index 96ace129f87..9f030d8b2d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx @@ -7,13 +7,27 @@ vi.mock("@/components/networking", () => ({ })); const mockMessagesStream = vi.fn(); +const mockMessagesCreate = vi.fn(); vi.mock("@anthropic-ai/sdk", () => ({ default: vi.fn(function () { - return { messages: { stream: mockMessagesStream } }; + return { messages: { stream: mockMessagesStream, create: mockMessagesCreate } }; }), })); +const NON_STREAMING_ARGS = [ + undefined, // traceId + undefined, // vector_store_ids + undefined, // guardrails + undefined, // policies + undefined, // selectedMCPServers + undefined, // customBaseUrl + undefined, // mcpServers + undefined, // mcpServerToolRestrictions + undefined, // mcpToolsets + false, // streamingEnabled +] as const; + describe("anthropic_messages prompt cache usage", () => { const captureUsage = async (usage: Record): Promise => { async function* mockStream() { @@ -59,3 +73,53 @@ describe("anthropic_messages prompt cache usage", () => { expect(usageData.promptTokens).toBe(5000); }); }); + +describe("anthropic_messages non-streaming", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("sends stream:false through messages.create and renders the full reply at once", async () => { + mockMessagesCreate.mockResolvedValue({ + content: [ + { type: "thinking", thinking: "considering" }, + { type: "text", text: "OK" }, + ], + usage: { input_tokens: 12, output_tokens: 3, cache_read_input_tokens: 7 }, + }); + const updateTextUI = vi.fn(); + const onReasoningContent = vi.fn(); + const onUsageData = vi.fn(); + + await makeAnthropicMessagesRequest( + [{ role: "user", content: "Hello" }], + updateTextUI, + "claude-haiku-4-5", + "test-token", + undefined, + undefined, + onReasoningContent, + undefined, + onUsageData, + ...NON_STREAMING_ARGS, + ); + + expect(mockMessagesStream).not.toHaveBeenCalled(); + expect(mockMessagesCreate).toHaveBeenCalledTimes(1); + expect(mockMessagesCreate.mock.calls[0][0]).toMatchObject({ model: "claude-haiku-4-5", stream: false }); + expect(updateTextUI).toHaveBeenCalledWith("assistant", "OK", "claude-haiku-4-5"); + expect(onReasoningContent).toHaveBeenCalledWith("considering"); + const expectedUsage: TokenUsage = { completionTokens: 3, promptTokens: 12, totalTokens: 15, cacheReadTokens: 7 }; + expect(onUsageData).toHaveBeenCalledWith(expectedUsage); + }); + + it("keeps streaming as the default when the flag is omitted", async () => { + async function* emptyStream() {} + mockMessagesStream.mockReturnValue(emptyStream()); + + await makeAnthropicMessagesRequest([{ role: "user", content: "Hello" }], vi.fn(), "claude-haiku-4-5", "test-token"); + + expect(mockMessagesCreate).not.toHaveBeenCalled(); + expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index 14afa013768..9dd2f675c44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -7,6 +7,13 @@ import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +const toTokenUsage = (usage: Anthropic.Usage): TokenUsage => ({ + completionTokens: usage.output_tokens, + promptTokens: usage.input_tokens, + totalTokens: usage.input_tokens + usage.output_tokens, + ...extractPromptCacheTokens(usage), +}); + export async function makeAnthropicMessagesRequest( messages: MessageType[], updateTextUI: (role: string, delta: string, model?: string) => void, @@ -26,6 +33,7 @@ export async function makeAnthropicMessagesRequest( mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, mcpToolsets?: MCPToolset[], + streamingEnabled: boolean = true, ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -58,7 +66,7 @@ export async function makeAnthropicMessagesRequest( const requestBody: any = { model: selectedModel, messages: messages.map((m) => ({ role: m.role, content: m.content })), - stream: true, + stream: streamingEnabled, max_tokens: 1024, // @ts-ignore - litellm specific parameter litellm_trace_id: traceId, @@ -74,6 +82,20 @@ export async function makeAnthropicMessagesRequest( if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; if (policies) requestBody.policies = policies; + + if (!streamingEnabled) { + const message: Anthropic.Message = await client.messages.create({ ...requestBody, stream: false }, { signal }); + for (const block of message.content) { + if (block.type === "text") { + updateTextUI("assistant", block.text, selectedModel); + } else if (block.type === "thinking" && onReasoningContent) { + onReasoningContent(block.thinking); + } + } + onUsageData?.(toTokenUsage(message.usage)); + return; + } + // Use the streaming helper method for cleaner async iteration // @ts-ignore - The SDK types might not include all litellm-specific parameters const stream = client.messages.stream(requestBody, { signal }); @@ -105,14 +127,7 @@ export async function makeAnthropicMessagesRequest( // Process usage data from message_delta events if (messageStreamEvent.type === "message_delta" && (messageStreamEvent as any).usage && onUsageData) { - const usage = (messageStreamEvent as any).usage; - const usageData: TokenUsage = { - completionTokens: usage.output_tokens, - promptTokens: usage.input_tokens, - totalTokens: usage.input_tokens + usage.output_tokens, - ...extractPromptCacheTokens(usage), - }; - onUsageData(usageData); + onUsageData(toTokenUsage((messageStreamEvent as any).usage)); } } } catch (error) { From ae6893bc6498d45785fad095821813d13e21cce3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:38:55 -0700 Subject: [PATCH 136/157] refactor(guardrails): narrow the Responses tool-call item mapping types --- .../openai/responses/guardrail_translation/handler.py | 4 ++-- .../responses/test_openai_responses_guardrail_handler.py | 8 +++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index e78e1f56915..fd0d7b66bee 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -193,11 +193,11 @@ def _is_tool_call_item(item: object) -> bool: return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES -def _tool_call_output_item_mapping(item: object) -> Mapping[str, Any] | None: +def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None: if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES: return None if isinstance(item, Mapping): - return cast("Mapping[str, Any]", item) # cast-ok: output items are str-keyed JSON objects + return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects return item.model_dump() if isinstance(item, BaseModel) else None diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 628ada05267..cf4895f637f 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -14,6 +14,7 @@ import pytest from fastapi import HTTPException +from pydantic import BaseModel from openai.types.responses import ( ResponseCustomToolCall, ResponseCustomToolCallInputDeltaEvent, @@ -28,6 +29,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -68,7 +70,7 @@ class PersimmonMaskingGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: Optional[LiteLLMLoggingObj] = None, ) -> GenericGuardrailAPIInputs: tool_calls = [ { @@ -593,7 +595,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction: texts_to_check: List[str] = [] images_to_check: List[str] = [] - tool_calls_to_check: List[Any] = [] + tool_calls_to_check: List[ChatCompletionToolCallChunk] = [] task_mappings: List[Tuple[int, int]] = [] # Extract tool calls @@ -1476,7 +1478,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: ) handler = OpenAIResponsesHandler() - typed_events: List[Any] = [ + typed_events: List[BaseModel] = [ model.model_validate({**event, "sequence_number": sequence_number}) for sequence_number, (model, event) in enumerate( zip( From eee239cd9e6b3b3e4bc3d3a515fa7d5160b51df9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 15:20:58 -0700 Subject: [PATCH 137/157] fix(jwt): cascade-delete JWT key mappings when their virtual key is deleted LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so deleting a virtual key that a JWT mapping pointed at failed with a foreign key violation on every deletion path (/key/delete, Admin UI, alias delete, team and user cascades). Declaring onDelete: Cascade on the relation lets the database clean the mapping up uniformly, so the next JWT call from that identity re-registers against the newly created key. Rebase of #33703 onto current staging. Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- .../migration.sql | 15 +++++ .../litellm_proxy_extras/schema.prisma | 2 +- litellm/proxy/schema.prisma | 2 +- schema.prisma | 2 +- .../test_litellm_proxy_extras_utils.py | 67 +++++++++++++++++++ 5 files changed, 85 insertions(+), 3 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql new file mode 100644 index 00000000000..e5d48abcb52 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql @@ -0,0 +1,15 @@ +-- DropForeignKey +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey"; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/schema.prisma b/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/schema.prisma +++ b/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 3fab20a28ad..e5826b18668 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -2,6 +2,7 @@ import glob import os import re import sys +from pathlib import Path import pytest @@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting: harness.run() assert len(harness.deploy_calls) == 1 assert harness.resolved == [] + + +class TestJWTKeyMappingCascade: + """Regression tests for issue #33702. + + A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted + because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so + deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign + key violation. The mapping must be removed automatically when its key is + deleted, which the FK now enforces via ON DELETE CASCADE. + """ + + _FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey" + + def _effective_on_delete(self): + """Replay every migration in order and return the last ON DELETE action + declared for the JWT key mapping FK.""" + action = None + for _migration_name, sql in _get_all_migrations(): + for match in re.finditer( + rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?' + r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)", + sql, + re.IGNORECASE | re.DOTALL, + ): + action = re.sub(r"\s+", " ", match.group(1).upper()) + return action + + def test_fk_effective_on_delete_is_cascade(self): + """The final FK definition across all migrations must cascade deletes.""" + assert self._effective_on_delete() == "CASCADE", ( + f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a " + "virtual key removes its JWT key mapping (issue #33702)" + ) + + def test_schema_declares_cascade_on_relation(self): + """schema.prisma must declare onDelete: Cascade on the mapping relation + so the generated client and DB agree.""" + schema_paths = glob.glob( + os.path.abspath( + os.path.join( + os.path.dirname(__file__), "../../**/schema.prisma" + ) + ), + recursive=True, + ) + declaring = tuple( + (path, schema) + for path, schema in ((p, Path(p).read_text()) for p in schema_paths) + if "model LiteLLM_JWTKeyMapping" in schema + ) + assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found" + for path, schema in declaring: + match = re.search( + r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)", + schema, + ) + assert match is not None, ( + f"{path} declares LiteLLM_JWTKeyMapping but its verification token " + "relation could not be parsed, so this test cannot vouch for it " + "(issue #33702)" + ) + assert "onDelete: Cascade" in match.group(1), ( + f"{path} must declare onDelete: Cascade on the JWT key mapping " + "relation (issue #33702)" + ) From 6e96885629521b69c01c56faa165300963f6fb9a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:40:15 -0700 Subject: [PATCH 138/157] fix(proxy): accept non-string callback vars in default_team_settings (#40458) * fix(proxy): accept non-string callback vars in default_team_settings A YAML boolean such as turn_off_message_logging: true in a default_team_settings block failed TeamCallbackMetadata's str-only callback_vars validation and errored the request before any callback ran. Stringify the value the same way AddTeamCallback does. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop docstring from default_team_settings bool regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): move default_team_settings bool regression test to mapped pre_call_utils suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 4 ++- .../proxy/test_litellm_pre_call_utils.py | 34 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2925226f235..da033dc2276 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1760,7 +1760,9 @@ class LiteLLMProxyRequestSetup: callback_vars_dict.pop("success_callback", None) callback_vars_dict.pop("failure_callback", None) callback_vars_dict = { - key: (litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else value) + key: ( + litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else str(value) + ) for key, value in callback_vars_dict.items() } 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 7564fa91c43..94aaa32519e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -35,6 +35,7 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_logging_from_dynamic_params from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -7909,3 +7910,36 @@ async def test_client_supplied_omit_marker_never_reaches_the_spend_log( if general_settings.get("missing_session_id") == "generate" else "per-call-random-trace-id" ) + + +def test_default_team_settings_bool_turn_off_message_logging_redacts(): + from litellm.proxy.proxy_server import ProxyConfig + + pc = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + { + "team_id": "team-redact", + "success_callback": ["gcs_bucket"], + "failure_callback": ["gcs_bucket"], + "turn_off_message_logging": True, + } + ] + } + } + + callback_metadata = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( + team_id="team-redact", + proxy_config=pc, + ) + + assert callback_metadata is not None + assert callback_metadata.success_callback == ["gcs_bucket"] + assert callback_metadata.callback_vars == {"turn_off_message_logging": "True"} + assert ( + _get_turn_off_message_logging_from_dynamic_params( + {"standard_callback_dynamic_params": dict(callback_metadata.callback_vars)} + ) + is True + ) From f5a410585af9fe2c06f0e5fa07ceddb2e0be46e0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:44:29 -0700 Subject: [PATCH 139/157] fix(responses): keep a custom tool_choice type in the bridged echo and type the new tests --- .../transformation.py | 23 +++++++++++++------ .../test_litellm_completion_responses.py | 17 +++++++++----- .../test_streaming_iterator_transformation.py | 9 ++++---- 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 878f493b58e..6b9931b7df2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,9 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -272,13 +273,21 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: - normalized: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) - match normalized: - case None: - return "auto" - case {"type": "function", "function": {"name": str(function_name)}}: + if tool_choice is None: + return "auto" + try: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice) + except ValidationError: + return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice) + + @staticmethod + def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice: + match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice): + case {"type": "custom"}, {"function": {"name": str(custom_name)}}: + return ToolChoiceCustomParam(type="custom", name=custom_name) + case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _: + case _, normalized: return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) @staticmethod diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2aab660b5b7..8f2937cb252 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,4 +1,5 @@ import json +from typing import Final import pytest @@ -1426,7 +1427,9 @@ class TestToolChoiceTransformation: [ ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), - ({"type": "custom", "name": "ApplyPatch"}, {"type": "function", "name": "ApplyPatch"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "function"}, "required"), ({"type": "tool"}, "required"), ({"type": "auto"}, "auto"), ("required", "required"), @@ -1434,14 +1437,16 @@ class TestToolChoiceTransformation: (None, "auto"), ], ) - def test_transform_tool_choice_for_responses_api_response(self, request_tool_choice, expected): - result = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + def test_transform_tool_choice_for_responses_api_response( + self, request_tool_choice: object, expected: str | dict[str, str] + ) -> None: + result: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( request_tool_choice ) assert result == expected - def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self): - chat_completion_response = ModelResponse( + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self) -> None: + chat_completion_response: Final = ModelResponse( id="chatcmpl-named-tool-choice", created=1748575031, model="claude-haiku-4-5", @@ -1465,7 +1470,7 @@ class TestToolChoiceTransformation: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( request_input="Run the command pwd.", responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, chat_completion_response=chat_completion_response, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index d4b565f82a1..29061db97dc 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -657,8 +658,8 @@ def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: ) -def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): - iterator = LiteLLMCompletionStreamingIterator( +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( model="claude-haiku-4-5", litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), request_input="Run the command pwd.", @@ -670,9 +671,9 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape(): litellm_metadata={}, ) - events = list(iterator) + events: Final = list(iterator) - response_events = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + response_events: Final = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] assert [event.type for event in response_events] == [ "response.created", "response.in_progress", From eb45a088d3d3c255118d0effb41e4613af0907c0 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 15:53:05 -0700 Subject: [PATCH 140/157] fix(router): resolve team-scoped auto-routers by their public name (#40432) A team-scoped auto-router is stored under an internal model_name_{team_id}_{uuid} with the caller-facing name in model_info.team_public_model_name, and the four pre-routing strategy registries key on that internal name. A team key asks for the public name, so the strategy lookup missed, the team early-resolve exit handed back the marker deployment itself, and every call 400'd with "Unmapped LLM provider". The strategy lookup now resolves the requested name through the same team-first, then global, then admin-across-teams deployment resolution the deployment path uses, and looks the registries up under the model_name of whatever that resolves to. Both exits of _common_checks_available_deployment drop strategy markers through one helper, so a marker-only resolution is rejected as uncallable on every path. The request team id has one reader. Resolves LIT-7363 Claude-Session: https://claude.ai/code/session_01NU97S7d2FUDDvTk59k53Wp Co-authored-by: Claude Fable 5.1 --- .../guardrails/auto_router_compression.py | 25 +- litellm/router.py | 144 +++++---- litellm/router_utils/common_utils.py | 14 +- .../test_auto_router_compression.py | 26 +- .../proxy/test_common_request_processing.py | 2 +- .../router_strategy/test_complexity_router.py | 4 +- tests/test_litellm/test_router.py | 301 +++++++++++++++++- 7 files changed, 417 insertions(+), 99 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 98707e7ddca..c37b9fff1f0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -80,17 +80,19 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( llm_router: "Router | None", model_alias: str, - team_id: str | None, + request_kwargs: Mapping[str, object], request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy of the auto router marker `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to for this caller. - Pre-call arming and the routing hook both resolve through here, so an alias with - several tag-scoped markers cannot suppress under one and then route under another. + Pre-call arming and the routing hook both resolve through here, and here resolves through the + router's own request-scoped deployment lookup, so an alias with several tag-scoped markers + cannot suppress under one and then route under another, and a team router reached by its + public name carries its policy for every principal that can reach it. """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs) markers: Final = tuple( litellm_params for deployment in deployments @@ -108,17 +110,6 @@ def policy_for_model( return next((policy for policy in candidates if policy is not None), None) -def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: - """The caller's team id, from whichever metadata bucket this surface writes to.""" - for meta_key in ("metadata", "litellm_metadata"): - meta = request_kwargs.get(meta_key) - if isinstance(meta, Mapping): - team_id = meta.get("user_api_key_team_id") - if isinstance(team_id, str): - return team_id - return None - - def _compression_guardrail_classes() -> tuple[type, ...]: """The registered guardrail classes whose provider compresses prompts.""" from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -172,7 +163,7 @@ async def arm_pre_call( policy: Final = policy_for_model( llm_router=llm_router, model_alias=model_alias, - team_id=team_id_from_request(data), + request_kwargs=data, request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: diff --git a/litellm/router.py b/litellm/router.py index 68ace283949..88601e5b97b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -148,6 +148,7 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + get_request_team_id, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -12337,27 +12338,7 @@ class Router: if team_deployments: return model, team_deployments elif include_team_models: - team_deployments = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) + team_deployments = self._team_deployments_across_teams(model) if team_deployments: return model, team_deployments @@ -12384,6 +12365,45 @@ class Router: return None + def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: + """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids: Final = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[DeploymentTypedDict]: + """The deployments `model` names for this caller, through the same alias, then team-first, then + global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so + strategy selection and compression policy can never disagree with deployment selection about + which marker a name means.""" + registered_name: Final = self._get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._get_all_deployments(model_name=registered_name, team_id=team_id) + if deployments or team_id is not None or not _is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered_name) + @staticmethod def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: litellm_params: Final = deployment.get("litellm_params") @@ -12411,11 +12431,7 @@ class Router: - Dict, if specific model chosen """ - request_team_id: str | None = None - if request_kwargs is not None: - metadata: Final = request_kwargs.get("metadata") or {} - litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) @@ -12440,7 +12456,9 @@ class Router: include_team_models=_is_proxy_admin_request(request_kwargs), ) if early is not None: - return early + if not isinstance(early[1], list): + return early + return early[0], self._drop_strategy_markers(early[0], early[1]) ## get healthy deployments ### get all deployments @@ -12517,19 +12535,22 @@ class Router: model ] # update the model to the actual value if an alias has been passed in - marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments) - if not any(marker_flags): - return model, healthy_deployments - selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters - d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker + return model, self._drop_strategy_markers(model, healthy_deployments) + + def _drop_strategy_markers( + self, model: str, deployments: Sequence[DeploymentTypedDict] + ) -> list[DeploymentTypedDict]: + """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" + selectable: Final = [ # mutable-ok: matches _common_checks_available_deployment's list contract + d for d in deployments if not self._is_strategy_marker_deployment(d) ] - if not selectable: + if deployments and not selectable: raise litellm.BadRequestError( message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", model=model, llm_provider="", ) - return model, selectable + return selectable def _filter_deployments_by_model_access_groups( self, @@ -13219,12 +13240,8 @@ class Router: return filtered - def _model_name_has_plain_deployments(self, model: str) -> bool: - indices: Final = self.model_name_to_deployment_indices.get(model) or () - return any(not self._is_strategy_marker_deployment(self.model_list[idx]) for idx in indices) - def _select_pre_routing_strategy( - self, model: str, request_kwargs: dict + self, model: str, request_kwargs: Mapping[str, object] ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": """ Resolve the pre-routing strategy for `model`, disambiguating deployments @@ -13235,6 +13252,12 @@ class Router: deployment the strategy was registered from via its (model_name, tags) pair. + The registries are keyed by the marker deployment's own `model_name`, which + for a team-scoped router is the internal `model_name_{team}_{uuid}` while + the caller sends the team's public name. So the names looked up are the + `model_name`s of whatever deployments this caller's request resolves `model` + to, and `model` itself when it resolves to none. + With tag filtering enabled, router-wide or by the request's enable_tag_filtering (which the proxy sets from key/team router_settings), strategies that all carry real tags matching none of @@ -13242,12 +13265,14 @@ class Router: deployments: returning None hands the request to ordinary tag-aware deployment selection. """ - candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [ - *self.auto_routers.get(model, []), - *self.complexity_routers.get(model, []), - *self.adaptive_routers.get(model, []), - *self.quality_routers.get(model, []), - ] + registries: Final = (self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers) + if not any(registries): + return None + deployments: Final = self.deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) + candidates: Final = tuple( + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) if not candidates: return None @@ -13265,7 +13290,7 @@ class Router: if ( (self.enable_tag_filtering or request_scoped_filtering) and all(tagged.tags for tagged in candidates) - and self._model_name_has_plain_deployments(model) + and any(not self._is_strategy_marker_deployment(d) for d in deployments) ): return None return candidates[0] @@ -13377,11 +13402,12 @@ class Router: Used for the litellm auto-router to modify the request before the routing decision is made. - `model` is whatever the caller asked for, which may be a `model_group_alias` key, while the - strategy registries and the marker deployment are keyed by the marker's own `model_name`, so - every lookup below resolves the alias first. Only the lookups: the caller-facing name stays - the alias, since spend metadata is stamped before routing and the response carries the tier - group the strategy picked. + `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's + public model name, while the strategy registries and the marker deployment are keyed by the + marker's own `model_name`, so every lookup below resolves the alias first and the team name + through the deployment path. Only the lookups: the caller-facing name stays the alias, since + spend metadata is stamped before routing and the response carries the tier group the + strategy picked. """ requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -13418,7 +13444,6 @@ class Router: messages_for_routing, model_hop_compression_armed, policy_for_model, - team_id_from_request, ) # Same tag-aware lookup the proxy's pre-call arming used, so an alias with @@ -13426,7 +13451,7 @@ class Router: compression_policy: Final = policy_for_model( llm_router=self, model_alias=registered_model_name, - team_id=team_id_from_request(request_kwargs), + request_kwargs=request_kwargs, request_tags=_get_tags_from_request_kwargs(request_kwargs), ) # Shared compression already ran in the pre-call hook, so reuse it rather than @@ -13495,7 +13520,9 @@ class Router: # Per-tier `litellm_params` on the hook response are deliberate overrides # the caller applies on top, so those keys are never forwarded here. marker_params: Final = ( - self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags) + self._forwardable_alias_marker_params( + model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs + ) if pre_routing_hook_response is not None else () ) @@ -13513,13 +13540,14 @@ class Router: return pre_routing_hook_response def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params - for idx in self.model_name_to_deployment_indices.get(model, ()) - if isinstance(litellm_params := self.model_list[idx].get("litellm_params", {}), dict) - and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + for deployment in self.deployments_for_request(model, request_kwargs) + if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith( + AUTO_ROUTER_MODEL_PREFIX + ) ) tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..c1c6d25beca 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -26,6 +26,18 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + if request_kwargs is None: + return None + for bucket_name in ("metadata", "litellm_metadata"): + bucket = request_kwargs.get(bucket_name) + team_id = bucket.get("user_api_key_team_id") if isinstance(bucket, Mapping) else None + if isinstance(team_id, str) and team_id: + return team_id + return None + + def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None: """ Resolve ``model`` through a ``model_group_alias`` map. @@ -110,7 +122,7 @@ def filter_team_based_models( metadata: Final = request_kwargs.get("metadata") or {} litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): requested_model: Final = ( request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index db2f94306fb..0d723d6671f 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -56,13 +56,13 @@ class TestPolicyFromLitellmParams: class _FakeRouter: - """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + """Minimal stand-in for litellm.Router.deployments_for_request, for policy_for_model.""" def __init__(self, deployments: list[dict[str, Any]]): self._deployments = deployments - def get_model_list(self, model_name, team_id=None): - return [d for d in self._deployments if d.get("model_name") == model_name] + def deployments_for_request(self, model, request_kwargs): + return [d for d in self._deployments if d.get("model_name") == model] def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: @@ -78,23 +78,23 @@ def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[ class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): router = _FakeRouter( [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_picks_the_marker_whose_tags_the_request_carries(self): @@ -107,8 +107,8 @@ class TestPolicyForModel: ] ) - eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) - us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + eu = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) @@ -116,7 +116,7 @@ class TestPolicyForModel: def test_untagged_marker_matches_any_request(self): router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) policy = policy_for_model( - llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("anything",) ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -128,14 +128,14 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-default"}), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): """No matching marker means no policy, not an unrelated slice's compression.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) assert ( - policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) is None ) def test_tag_scoped_marker_takes_precedence_over_untagged(self): @@ -147,7 +147,7 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index d2ba6527826..50b26577e5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -417,7 +417,7 @@ class TestProxyBaseLLMRequestProcessing: ) fake_llm_router = MagicMock() - fake_llm_router.get_model_list.return_value = [ + fake_llm_router.deployments_for_request.return_value = [ { "model_name": "smart-router", "litellm_params": { diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 9f925499f1a..51103297c58 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3846,11 +3846,11 @@ class TestRouterPreRoutingSharedAliasName: def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self): router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()]) - forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=())) + forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={})) assert forwarded["drop_params"] is True assert "api_key" not in forwarded and "api_base" not in forwarded - assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=()) == () + assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == () @staticmethod def _region_marker_entry() -> dict: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 80da922724d..1e0327e5ac2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10003,13 +10003,6 @@ class TestTaggedAutoRouterOnSharedModelName: def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - def test_model_name_has_plain_deployments_reflects_the_pool(self): - mixed = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True) - marker_only = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True) - - assert mixed._model_name_has_plain_deployments("gpt4o") is True - assert marker_only._model_name_has_plain_deployments("gpt4o") is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -10618,6 +10611,300 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestTeamPublicNameReachesPreRoutingStrategies: + """A team-scoped strategy router is stored under an internal `model_name_{team}_{uuid}` with the + caller-facing name in `model_info.team_public_model_name`, and the four registries key on that + internal name. A team key asks for the public name, so the hook has to resolve it to the team's + marker through the same team-first resolution the deployment path uses, and a resolution that + yields only markers is not callable on any path (LIT-7363).""" + + MARKER_TIMEOUT = 42.0 + REGISTRY_NAMES = ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers") + TEAM = "team-a" + OTHER_TEAM = "team-b" + PUBLIC_NAME = "smart-route" + INTERNAL_NAME = "model_name_team-a_0b3c" + SIBLING_INTERNAL_NAME = "model_name_team-a_9e1d" + + class _RewriteStrategy: + def __init__(self, rewrite_to: str = "gemini-flash"): + self.rewrite_to = rewrite_to + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse(model=self.rewrite_to, messages=messages) + + @classmethod + def _team_marker(cls, internal_name: str, tags: list[str] | None = None) -> dict: + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + return { + "model_name": internal_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + "timeout": cls.MARKER_TIMEOUT, + **({"tags": tags} if tags else {}), + }, + "model_info": {"team_id": cls.TEAM, "team_public_model_name": cls.PUBLIC_NAME}, + } + + @classmethod + def _router( + cls, + registrations: dict[str, "TestTeamPublicNameReachesPreRoutingStrategies._RewriteStrategy"], + registry_name: str = "complexity_routers", + extra_deployments: tuple[dict, ...] = (), + markers: tuple[dict, ...] | None = None, + enable_tag_filtering: bool = False, + ) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + markers = markers if markers is not None else (cls._team_marker(cls.INTERNAL_NAME),) + tier = { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + } + router = litellm.Router( + model_list=[*markers, tier, *extra_deployments], + enable_tag_filtering=enable_tag_filtering, + ) + tags_by_name = {m["model_name"]: tuple(m["litellm_params"].get("tags") or ()) for m in markers} + for name in cls.REGISTRY_NAMES: + setattr(router, name, {}) + setattr( + router, + registry_name, + { + name: [TaggedPreRoutingStrategy(tags=tags_by_name[name], strategy=strategy)] + for name, strategy in registrations.items() + }, + ) + return router + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + @classmethod + def _team_request(cls, team_id: str | None = "team-a", tags: list[str] | None = None) -> dict: + metadata = {**({"user_api_key_team_id": team_id} if team_id else {}), **({"tags": tags} if tags else {})} + return {"metadata": metadata} + + @pytest.mark.parametrize("registry_name", REGISTRY_NAMES) + @pytest.mark.asyncio + async def test_team_key_dispatches_to_the_strategy_registered_under_the_internal_name(self, registry_name): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}, registry_name=registry_name) + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + + @pytest.mark.asyncio + async def test_team_key_deployment_selection_lands_on_the_tier_and_forwards_the_marker_params(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = self._team_request() + + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert deployment["litellm_params"]["model"] == "gemini/gemini-3.6-flash" + assert request_kwargs["timeout"] == self.MARKER_TIMEOUT + + @pytest.mark.asyncio + async def test_another_team_never_reaches_the_strategy_or_the_marker(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + + assert ( + await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + is None + ) + with pytest.raises(litellm.BadRequestError): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_sibling_team_markers_select_by_request_tag_then_default(self): + router = self._router( + { + self.INTERNAL_NAME: self._RewriteStrategy("cn-model"), + self.SIBLING_INTERNAL_NAME: self._RewriteStrategy("us-model"), + }, + markers=( + self._team_marker(self.INTERNAL_NAME, tags=["cn"]), + self._team_marker(self.SIBLING_INTERNAL_NAME, tags=["us", "default"]), + ), + ) + + async def routed(tags: list[str] | None) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=tags), messages=self._messages() + ) + return response.model if response else None + + assert await routed(["cn"]) == "cn-model" + assert await routed(["us"]) == "us-model" + assert await routed(None) == "us-model" + + @pytest.mark.asyncio + async def test_team_public_name_shadows_a_global_model_for_that_team_only(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": self.PUBLIC_NAME, "litellm_params": {"model": "openai/gpt-4o"}},), + ) + + async def routed(request_kwargs: dict) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return response.model if response else None + + async def selected(request_kwargs: dict) -> str: + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return deployment["litellm_params"]["model"] + + assert await routed(self._team_request()) == "gemini-flash" + assert await selected(self._team_request()) == "gemini/gemini-3.6-flash" + for request_kwargs in (self._team_request(None), self._team_request(self.OTHER_TEAM)): + assert await routed(request_kwargs) is None + assert await selected(request_kwargs) == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_tag_filtering_hands_untagged_team_requests_to_the_team_plain_sibling(self): + plain_sibling = { + "model_name": self.SIBLING_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"team_id": self.TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + markers=(self._team_marker(self.INTERNAL_NAME, tags=["route"]),), + extra_deployments=(plain_sibling,), + enable_tag_filtering=True, + ) + + tagged = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=["route"]), messages=self._messages() + ) + assert tagged is not None and tagged.model == "gemini-flash" + for _ in range(20): + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + assert deployment["litellm_params"]["model"] == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_marker_only_team_resolution_is_rejected_as_uncallable(self): + import re + + from litellm.types.router import RouterErrors + + router = self._router({}) + + with pytest.raises( + litellm.BadRequestError, match=re.escape(RouterErrors.only_strategy_marker_deployments.value) + ): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_proxy_admin_without_a_team_reaches_the_team_strategy_by_public_name(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + @pytest.mark.asyncio + async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), + ) + principals = { + "team": self._team_request(), + "other-team": self._team_request(self.OTHER_TEAM), + "no-team": self._team_request(None), + "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, + } + for principal, request_kwargs in principals.items(): + for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): + resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] + callable_names = [ + name + for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) + if not router._is_strategy_marker_deployment(deployment) + ] + if resolved and not callable_names: + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + elif not resolved: + with pytest.raises(litellm.BadRequestError): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + else: + _, deployments = router._common_checks_available_deployment( + model=model, request_kwargs=request_kwargs + ) + assert [d["model_name"] for d in deployments] == callable_names, (principal, model) + + def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): + router = self._router({}) + marker = router.model_list[0] + plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} + + assert router._drop_strategy_markers("x", [marker, plain]) == [plain] + assert router._drop_strategy_markers("x", [plain]) == [plain] + assert router._drop_strategy_markers("x", []) == [] + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._drop_strategy_markers("x", [marker]) + + def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): + other_team_marker = { + **self._team_marker(self.SIBLING_INTERNAL_NAME), + "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + one_team = self._router({}) + two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) + + assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ + self.INTERNAL_NAME + ] + assert one_team._team_deployments_across_teams("missing") == [] + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + two_teams._team_deployments_across_teams(self.PUBLIC_NAME) + + + def test_compression_policy_follows_the_same_resolution_for_every_principal(self): + from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model + + marker = self._team_marker(self.INTERNAL_NAME) + marker["litellm_params"]["auto_router_routing_compression"] = "headroom-team" + router = self._router({}, markers=(marker,)) + admin = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + expected = AutoRouterCompressionPolicy(routing="headroom-team", model=None) + + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(), ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, admin, ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(self.OTHER_TEAM), ()) is None + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(None), ()) is None + + class TestAutoRouterCompressionDecoupling: """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` decouple what the routing decision sees from what the model call sees. The one From fc161faa96a27f16d90d1aee6e4fc0574f679dd9 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 9 Sep 2026 15:56:38 -0700 Subject: [PATCH 141/157] fix(langfuse): give each call in a session header its own trace instead of upserting one trace per session (#40177) * fix(langfuse): give session-header calls their own trace id A client that sends only a session header (x-litellm-session-id, a vendor x--session-id such as Claude Code's X-Claude-Code-Session-Id, a bare x-session-id, or the Codex session/thread/conversation family) has that value stamped into both trace_id and session_id by the proxy. Langfuse upserts a trace by id, so every turn of a session collapsed into one growing trace and the Sessions view showed "Total traces: 1" Detect that aliasing in the Langfuse callback from the request headers the callback already receives, and use litellm_call_id as the trace id for those calls. session_id still carries the header value, so the turns stay grouped under one session. An explicit x-litellm-trace-id, langfuse_trace_id, or langfuse_existing_trace_id keeps its trace id, including when the caller sets it to the same value as the session id Co-authored-by: jesus * test(langfuse): cover direct-SDK callers without proxy request headers * fix(langfuse): preserve session trace provenance --------- Co-authored-by: jesus --- litellm/integrations/langfuse/langfuse.py | 59 +++- .../integrations/test_langfuse.py | 251 ++++++++++++++++++ 2 files changed, 306 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 9576eabaa34..b75369965de 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -2,6 +2,7 @@ # On success, logs events to Langfuse import inspect import os +import re import traceback from collections.abc import Callable, Iterable, Mapping from datetime import datetime @@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, dict) else None +def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]: + """Header pairs with the key type widened back to what a caller-supplied dict can actually hold.""" + return mapping.items() + + +def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool: + if not isinstance(trace_id, str) or not isinstance(session_id, str): + return False + request: Final = _object_mapping(proxy_server_request) + raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None + if raw_headers is None: + return False + headers: Final = MappingProxyType( + {key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)} + ) + if headers.get("x-litellm-trace-id"): + return False + if headers.get("langfuse_trace_id") is not None: + return False + if trace_id != session_id and headers.get("langfuse_session_id") != session_id: + return False + if headers.get("x-litellm-session-id") == trace_id: + return True + if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None: + return False + user_agent: Final = headers.get("user-agent") + codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None + return any( + value == trace_id + and ( + key == "x-session-id" + or re.fullmatch(r"x-.+-session-id", key) is not None + or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id")) + ) + for key, value in headers.items() + ) + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -609,6 +648,18 @@ class LangFuseLogger: # This allows continuing an existing trace while still returning the correct trace_id if existing_trace_id is not None: trace_id = existing_trace_id + resolved_trace_id: Final = ( + litellm_call_id or trace_id + if existing_trace_id is None + and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request")) + else trace_id + ) + if resolved_trace_id != trace_id: + verbose_logger.debug( + "Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace", + trace_id, + resolved_trace_id, + ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () @@ -663,7 +714,7 @@ class LangFuseLogger: trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { - "id": trace_id, + "id": resolved_trace_id, "name": trace_name, "session_id": session_id, "input": masked_input if not mask_input else "redacted-by-litellm", @@ -845,13 +896,13 @@ class LangFuseLogger: # Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value # to match expected test behavior if hasattr(generation_client, "trace_id") and generation_client.trace_id: - if generation_client.trace_id != trace_id: + if generation_client.trace_id != resolved_trace_id: verbose_logger.warning( "Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.", - trace_id, + resolved_trace_id, generation_client.trace_id, ) - return trace_id, generation_id + return resolved_trace_id, generation_id except Exception: verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc()) return None, None diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index d36878e455f..87e76499b84 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1341,6 +1341,257 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None): ) +@pytest.mark.parametrize("level", ["DEFAULT", "ERROR"]) +@pytest.mark.parametrize( + "headers,metadata,expected_id", + [ + ({"x-litellm-session-id": "session-7125"}, {}, "call"), + ({"X-Claude-Code-Session-Id": "session-7125"}, {}, "call"), + ({"x-session-id": "session-7125"}, {}, "call"), + ({"session-id": "session-7125", "user-agent": "codex_cli_rs/1.0"}, {}, "call"), + ({"thread-id": "session-7125", "user-agent": "codex-tui"}, {}, "call"), + ({"session_id": "session-7125", "user-agent": "Codex 1.0"}, {}, "call"), + ({"conversation_id": "session-7125", "user-agent": "codex_vscode/1.0"}, {}, "call"), + ({"x-litellm-session-id": "short"}, {}, "call"), + ({"x-litellm-trace-id": "session-7125"}, {}, "session-7125"), + ( + {"X-LiteLLM-Trace-Id": "session-7125", "x-litellm-session-id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "explicit-trace"}, + {}, + "explicit-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_existing_trace_id": "existing-trace"}, + {}, + "existing-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-litellm-session-id": "short", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"X-Claude-Code-Session-Id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + { + "session-id": "session-7125", + "user-agent": "codex_cli_rs/1.0", + "langfuse_session_id": "custom-session", + }, + {}, + "call", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "x-litellm-trace-id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_trace_id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_existing_trace_id": "existing-trace", + }, + {}, + "existing-trace", + ), + ({}, {"trace_id": "session-7125", "session_id": "session-7125"}, "session-7125"), + ({}, {"trace_id": "explicit-trace", "session_id": "session-7125"}, "explicit-trace"), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "short", "session_id": "short"}, + "short", + ), + ( + {"x-session-id": "invalid value"}, + {"trace_id": "invalid value", "session_id": "invalid value"}, + "invalid value", + ), + ( + {"session-id": "session-7125", "user-agent": "codexfoo/1.0"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ({}, {}, "call"), + ], +) +def test_session_header_trace_provenance(headers, metadata, expected_id, level): + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + redact_credential_headers, + ) + + logger: Final = _steering_logger() + for turn in range(2): + call_id = f"call-{turn}" + request_headers = Headers(headers) + data = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=request_headers, data={"metadata": dict(metadata)}, _metadata_variable_name="metadata" + ) + original_metadata = dict(data["metadata"]) + now = datetime.datetime.now() + result = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": call_id, + "litellm_trace_id": data.get("litellm_trace_id"), + "litellm_params": { + "metadata": data["metadata"], + "proxy_server_request": {"headers": redact_credential_headers(clean_headers(request_headers))}, + }, + "messages": [{"role": "user", "content": f"turn {turn}"}], + "optional_params": {}, + }, + response_obj=( + None + if level == "ERROR" + else litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]) + ), + start_time=now, + end_time=now, + level=level, + status_message="provider error" if level == "ERROR" else None, + ) + trace_params = logger.Langfuse.trace.call_args.kwargs + assert trace_params["id"] == (call_id if expected_id == "call" else expected_id) + assert result["trace_id"] == trace_params["id"] + if expected_id != "existing-trace": + assert trace_params["session_id"] == headers.get("langfuse_session_id", original_metadata.get("session_id")) + steering = {key[len("langfuse_") :]: value for key, value in headers.items() if key.startswith("langfuse_")} + assert data["metadata"] == {**original_metadata, **steering} + + +def test_session_header_trace_without_call_id_keeps_session_alias(): + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": {"headers": {"x-litellm-session-id": "session-7125"}}, + }, + "messages": [{"role": "user", "content": "no call id"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_every_proxy_session_header_shape_is_classified_as_a_session_alias(): + """The classifier must cover every header shape the proxy turns into a chain id.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + from litellm.proxy.litellm_pre_call_utils import ( + _CODEX_SESSION_ID_HEADERS, + get_chain_id_from_headers, + ) + + session: Final = "session-7125-abcdef" + session_shapes: Final = ( + {"x-litellm-session-id": session}, + {"X-Claude-Code-Session-Id": session}, + {"x-session-id": session}, + *({header: session, "user-agent": "codex_cli_rs/1.0"} for header in _CODEX_SESSION_ID_HEADERS), + ) + for headers in session_shapes: + assert get_chain_id_from_headers(dict(headers)) == session, headers + assert _is_session_header_trace(session, session, {"headers": headers}) is True, headers + + explicit_trace: Final = {"x-litellm-trace-id": session, "x-litellm-session-id": session} + assert get_chain_id_from_headers(dict(explicit_trace)) == session + assert _is_session_header_trace(session, session, {"headers": explicit_trace}) is False + + +@pytest.mark.parametrize( + "proxy_server_request", + [None, {}, {"headers": None}], + ids=["no-proxy-request", "no-headers-key", "null-headers"], +) +def test_sdk_caller_without_request_headers_keeps_its_trace(proxy_server_request): + """A direct SDK caller has no request headers, so a session-shaped trace id stays the caller's.""" + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "call-0", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": proxy_server_request, + }, + "messages": [{"role": "user", "content": "sdk turn"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_session_header_classifier_survives_non_string_header_keys(): + """A non-string header key must not cost the caller its whole trace.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + + session: Final = "session-7125-abcdef" + headers: Final = {7: "numeric key", "x-litellm-session-id": session} + assert _is_session_header_trace(session, session, {"headers": headers}) is True + assert _is_session_header_trace(session, session, {"headers": {7: "numeric key"}}) is False + + def test_mask_input_header_false_keeps_the_prompt(): logger = _steering_logger() From 21e6c6e2f3c9fe993925b6aaaca6314bd56b1add Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:56:41 -0700 Subject: [PATCH 142/157] test(convert_dict_to_response): expect the type-naming error for a null choices value --- .../test_convert_dict_to_chat_completion.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 09d1adba17c..31c554985a7 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1645,7 +1645,7 @@ class TestMissingChoicesGuard: assert result.usage.prompt_tokens == 10 def test_convert_to_model_response_object_null_choices_raises_api_error(self): - """choices=None raises APIError.""" + """choices=None raises APIError that names the type instead of claiming the key is missing.""" from litellm.exceptions import APIError response_object = { @@ -1661,7 +1661,7 @@ class TestMissingChoicesGuard: model_response_object=ModelResponse(), ) - assert "no 'choices'" in exc_info.value.message + assert "'choices' that is not a list (NoneType)" in exc_info.value.message def test_convert_to_streaming_response_no_choices_raises_api_error(self): """Missing choices in streaming cache-hit path raises APIError.""" From fbc6fb56ae2d8d8e600b0e80f7a46067b796d47f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:06:19 -0700 Subject: [PATCH 143/157] fix(bedrock): route gpt-6-astra reasoning_effort to reasoning.effort and mark Nova 2 tool_choice The converse reasoning gate only matched openai.gpt-5, so gpt-6-astra fell through to Anthropic's thinking block and Bedrock rejected the call with 400 Unknown parameter: 'thinking'. Match any openai.gpt- model at the three gate sites instead. Nova 2 lite and pro accept forced tool_choice on Converse (verified live on us.amazon.nova-2-lite-v1:0), so the nine Nova 2 registry keys now advertise supports_tool_choice. The invoke dispatcher also forwards json_mode to Nova like it already does for Anthropic and TwelveLabs. --- .../bedrock/chat/converse_transformation.py | 17 +++++-- .../base_invoke_transformation.py | 1 + ...odel_prices_and_context_window_backup.json | 9 ++++ model_prices_and_context_window.json | 9 ++++ .../test_base_invoke_transformation.py | 44 +++++++++++++++++++ .../chat/test_converse_transformation.py | 3 ++ 6 files changed, 79 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 6173f08fc87..4cdf387a11d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse` import copy import json +import re import time import types from collections.abc import Mapping @@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) + @staticmethod + def _is_openai_gpt_reasoning_model(model: str) -> bool: + return re.search(r"openai\.gpt-\d", model) is not None + def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -423,14 +428,14 @@ class AmazonConverseConfig(BaseConfig): Handle the reasoning_effort parameter based on the model type. - GPT-OSS models: passed through unchanged via additionalModelRequestFields. - - OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields. + - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ if "gpt-oss" in model: optional_params["reasoning_effort"] = reasoning_effort - elif "openai.gpt-5" in model: + elif self._is_openai_gpt_reasoning_model(model): reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} optional_params["reasoning"] = reasoning elif self._is_nova_2_model(model): @@ -564,7 +569,11 @@ class AmazonConverseConfig(BaseConfig): # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html supported_params.append("tool_choice") - if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model: + if ( + "gpt-oss" in model + or self._is_openai_gpt_reasoning_model(model) + or self._is_openai_gpt_reasoning_model(base_model) + ): supported_params.append("reasoning_effort") elif self._is_nova_2_model(model): # Nova 2 models support reasoning_effort (transformed to reasoningConfig) @@ -920,7 +929,7 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } - if param == "thinking" and "openai.gpt-5" not in model: + if param == "thinking" and not self._is_openai_gpt_reasoning_model(model): if ( isinstance(value, dict) and value.get("type") == "adaptive" diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index a0e32c8aa22..90a2692f68a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, + json_mode=json_mode, ) elif provider == "twelvelabs": return litellm.AmazonTwelveLabsPegasusConfig().transform_response( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a2e7b692649..0d2eda93323 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -381,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -400,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -417,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -436,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -453,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -472,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -489,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -508,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -28578,6 +28586,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a2e7b692649..0d2eda93323 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -381,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -400,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -417,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -436,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -453,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -472,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -489,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -508,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -28578,6 +28586,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index c2c448cd7e2..96a2fa6ec67 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -1,5 +1,7 @@ import json +from unittest.mock import MagicMock +import httpx import pytest @@ -190,3 +192,45 @@ def test_get_error_class_preserves_provider_headers(): assert isinstance(error, BedrockError) assert error.headers == {"x-amzn-RequestId": "req-invoke-500"} assert error.response.headers["x-amzn-requestid"] == "req-invoke-500" + + +def test_transform_response_hands_json_mode_to_nova(): + """The invoke dispatcher forwards its json_mode argument to Nova instead of dropping it.""" + from litellm.types.utils import ModelResponse + + response_json = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_nova_json", + "name": "json_tool_call", + "input": {"city": "Paris", "temperature": 21}, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock")) + + result = AmazonInvokeConfig().transform_response( + model="invoke/amazon.nova-lite-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "weather"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=True, + ) + + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 16824d19f23..51c0cd261e8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -382,6 +382,8 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): "us.openai.gpt-5.6-sol", "global.openai.gpt-5.6-terra", "bedrock/converse/us.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", + "bedrock/converse/global.openai.gpt-6-astra", ], ) def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(model, local_model_cost_map): @@ -412,6 +414,7 @@ def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(mode [ "us.openai.gpt-5.6-sol", "bedrock/converse/global.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", ], ) def test_openai_gpt5_converse_never_forwards_thinking(model, local_model_cost_map): From 96e46ba8ff43cdbf24d53691685c15b533bfa746 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:06:55 -0700 Subject: [PATCH 144/157] fix(responses): stamp streamed usage cost when the provider usage arrives as a dict Perplexity's Responses payload fails ResponsesAPIResponse validation on truncation "" and is kept as an unvalidated model, so its usage stays a plain dict and _stamp_responses_usage_cost raised AttributeError on every streamed completion once reasoning made the cost non-zero. Validate the dict into ResponseAPIUsage before stamping, keeping a provider-reported cost when it carries one. Resolves LIT-7391 --- litellm/responses/streaming_iterator.py | 17 ++++++- .../responses/test_streaming_iterator.py | 51 +++++++++++++++++++ 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..658a0e8ad5d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder +from pydantic import ValidationError from typing_extensions import TypeIs import litellm @@ -544,7 +545,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return try: @@ -1293,14 +1294,26 @@ def _add_text_like_part_events( ) +def _usage_as_model(usage: object) -> ResponseAPIUsage | None: + if isinstance(usage, ResponseAPIUsage): + return usage + if not isinstance(usage, dict): + return None + try: + return ResponseAPIUsage.model_validate(usage) + except ValidationError: + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: if response_obj is None or logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return + response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives if isinstance(getattr(usage_obj, "cost", None), (int, float)): return try: diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c226c0b4d09..60ce08666b6 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -368,6 +368,57 @@ def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): logging_obj._response_cost_calculator.assert_not_called() +def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: + return ResponsesAPIResponse.model_construct( + id="resp_lit7391", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="perplexity/deepseek-v4-flash-0731", + object="response", + output=[], + truncation="", + usage=usage, + ) + + +def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.types.llms.openai import ResponseAPIUsage + + response = _unvalidated_response_with_dict_usage( + { + "input_tokens": 29, + "output_tokens": 120, + "output_tokens_details": {"reasoning_tokens": 117}, + "total_tokens": 149, + "cost": {"currency": "USD", "input_cost": 0, "output_cost": 3e-05, "total_cost": 3e-05}, + } + ) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(3e-05) + assert response.usage.output_tokens_details.reasoning_tokens == 117 + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.types.llms.openai import ResponseAPIUsage + + response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + def test_stamp_responses_usage_cost_survives_calculator_failure(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost From a248ff2c6ea712b0f8cb9f03bc327701c8dd63ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:14:50 -0700 Subject: [PATCH 145/157] fix(guardrails): fail open on Responses tool-call rewrites in another shape and keep item names Guardrails that hand back tool_calls in their own shape (vendor JSON, user code output) raised a KeyError on the non-stream Responses write-back. Returned tool calls are now validated before comparison; a shape or count that does not line up leaves every tool-call item unchanged and logs a warning naming the guardrail. A tool call's name is written back only when the guardrail changed it, so a nameless custom_tool_call no longer picks up the custom_tool placeholder. --- .../guardrail_translation/handler.py | 85 ++++++++++++++---- ...test_openai_responses_guardrail_handler.py | 90 +++++++++++++++++++ 2 files changed, 159 insertions(+), 16 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index fd0d7b66bee..8c8294ad2fc 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -37,7 +37,7 @@ from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -105,6 +105,19 @@ class _ToolCallShape(NamedTuple): arguments: str +class _ToolCallFunctionFields(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + arguments: str = "" + + +class _ToolCallFields(BaseModel): + model_config = ConfigDict(frozen=True) + + function: _ToolCallFunctionFields + + def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: return tuple( _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) @@ -112,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl ) +def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None: + payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + try: + fields: Final = _ToolCallFields.model_validate(payload) + except ValidationError: + return None + return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments) + + +def _post_guardrail_tool_call_shapes( + returned_tool_calls: Sequence[object] | None, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str | None, +) -> tuple[_ToolCallShape, ...]: + if not pre_guardrail_tool_calls: + return pre_guardrail_tool_calls + if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, " + "leaving the tool call output items unchanged", + guardrail_name, + "no" if returned_tool_calls is None else len(returned_tool_calls), + len(pre_guardrail_tool_calls), + ) + return pre_guardrail_tool_calls + returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls) + validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None) + if len(validated_shapes) != len(returned_shapes): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, " + "leaving the tool call output items unchanged", + guardrail_name, + ) + return pre_guardrail_tool_calls + return validated_shapes + + +def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape: + return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -698,7 +752,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -709,11 +767,7 @@ class OpenAIResponsesHandler(BaseTranslation): self._write_tool_call_rewrites_to_output( tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)), pre_guardrail_tool_calls=pre_guardrail_tool_calls, - post_guardrail_tool_calls=_tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check - ), + post_guardrail_tool_calls=post_guardrail_tool_calls, ) verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) @@ -811,11 +865,10 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") - post_guardrail_tool_calls: Final = _tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, ) # Write guardrailed texts back into the output items in-place. @@ -991,7 +1044,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) rewrites_by_call_id: Final = MappingProxyType( { - call_id: after + call_id: _tool_call_rewrite(before, after) for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before } @@ -1050,12 +1103,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) -> None: if len(tool_call_items) != len(post_guardrail_tool_calls): return - for output_item, after in ( - (output_item, after) + for output_item, rewrite in ( + (output_item, _tool_call_rewrite(before, after)) for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before ): - self._write_tool_call_item(output_item, after.name, after.arguments) + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) @staticmethod def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index cf4895f637f..a4f0a77a9b6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -10,6 +10,8 @@ from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock +import logging + import pytest @@ -85,6 +87,29 @@ class PersimmonMaskingGuardrail(CustomGuardrail): return {**inputs, "tool_calls": tool_calls} +class FlatShapeGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])] + return {**inputs, "tool_calls": flat_tool_calls} + + +class DroppingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tool_calls": []} + + CUSTOM_TOOL_CALL_ITEM = { "type": "custom_tool_call", "id": "ctc_1", @@ -737,6 +762,52 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert (custom_item.name if typed else custom_item["name"]) == "exec" assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon" + @staticmethod + def _custom_tool_call_response(item: dict) -> dict: + return { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [item], + } + + @pytest.mark.asyncio + async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat")) + + assert result["output"][0]["input"] == "echo persimmon" + assert result["output"][0]["name"] == "exec" + + @pytest.mark.asyncio + async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper")) + + assert result["output"][0]["input"] == "echo persimmon" + assert any( + "dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"} + response = self._custom_tool_call_response(nameless_item) + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + assert result["output"][0]["input"] == "echo [MASKED]" + assert "name" not in result["output"][0] + @pytest.mark.asyncio async def test_process_output_response_with_tool_calls(self): """Test processing output response containing function tool calls""" @@ -1469,6 +1540,25 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[5]["response"]["output"][0]["name"] == "exec" assert "arguments" not in events[5]["response"]["output"][0] + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]] + for item in items: + del item["name"] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[3]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert all("name" not in item for item in items) + @pytest.mark.asyncio async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self): from litellm.types.llms.openai import ( From 7563d94e65591c2ed3b909c196a3eb76e61625fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:29:55 -0700 Subject: [PATCH 146/157] test(responses): assert the streamed response.completed event echoes the named tool_choice --- .../test_streaming_iterator_transformation.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 29061db97dc..70f0957890d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -679,6 +679,9 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: "response.in_progress", "response.completed", ] - assert response_events[0].response.tool_choice == {"type": "function", "name": "run_command"} - assert response_events[1].response.tool_choice == {"type": "function", "name": "run_command"} + assert [event.response.tool_choice for event in response_events] == [ + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + ] assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) From a68fe4e4d9eb8b55efd33be65ef4e1f1f1932c62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:31:48 -0700 Subject: [PATCH 147/157] fix(responses): keep the client's usage shape when the logging copy cannot re-validate the response --- litellm/responses/streaming_iterator.py | 41 ++++++++------ .../responses/test_streaming_iterator.py | 54 ++++++++++++++++--- 2 files changed, 72 insertions(+), 23 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 658a0e8ad5d..40ff88fc557 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from typing_extensions import TypeIs import litellm @@ -439,18 +439,7 @@ class BaseResponsesAPIStreamingIterator: if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): - try: - logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) - except Exception: - # Fallback to original if serialization fails - pass + logging_response: Final[object] = _logging_copy(self.completed_response) self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() @@ -489,10 +478,10 @@ class BaseResponsesAPIStreamingIterator: def _restore_provider_response_headers(self, logging_response: object) -> None: """Re-apply the provider's response headers to the copy handed to logging callbacks. - ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the ``_hidden_params`` the provider transform set on the nested response are lost. Returns early - when that copy fell back to the original event, so logging-only state never lands on the - object the caller is iterating. + when the event was not a pydantic model and logging got the original, so logging-only state + never lands on the object the caller is iterating. """ if logging_response is self.completed_response: return @@ -1294,6 +1283,26 @@ def _add_text_like_part_events( ) +def _logging_copy(event: object) -> object: + """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never + reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the + deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow + copies of the event and its nested response still keep the caller's ``usage`` attribute separate.""" + if not isinstance(event, BaseModel): + return event + try: + return type(event).model_validate(event.model_dump()) + except Exception: + return _detached_shallow_copy(event) + + +def _detached_shallow_copy(event: BaseModel) -> BaseModel: + nested: Final[object] = getattr(event, "response", None) + if isinstance(nested, BaseModel): + return event.model_copy(update={"response": nested.model_copy()}) + return event.model_copy() + + def _usage_as_model(usage: object) -> ResponseAPIUsage | None: if isinstance(usage, ResponseAPIUsage): return usage diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 60ce08666b6..5e0e794d93e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import ( SyncResponsesAPIStreamingIterator, ) from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): def _responses_api_response_with_usage() -> ResponsesAPIResponse: - from litellm.types.llms.openai import ResponseAPIUsage - return ResponsesAPIResponse( id="resp_lit6427", created_at=int(datetime(2025, 1, 1).timestamp()), @@ -383,8 +382,6 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage( { "input_tokens": 29, @@ -406,8 +403,6 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj._response_cost_calculator.return_value = 0.000704 @@ -586,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): iterator._log_completed_response(is_async=True) - assert logged == [iterator.completed_response] + assert len(logged) == 1 + assert logged[0] is not iterator.completed_response + assert logged[0].response is not iterator.completed_response.response + assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" assert iterator.completed_response.response._hidden_params == {} + + +def _unvalidated_completed_config() -> Mock: + """Config whose completed event carries a Perplexity-style response that fails validation + (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + response = _unvalidated_response_with_dict_usage( + ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001}) + ) + return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation(): + """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging + rewrites the assembled response's usage to chat shape in place, so the event handed to logging + must never be the one the caller receives.""" + logging_obj = _logging_obj_stub() + logging_obj.stream = True + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj) + events = [event async for event in iterator] + + assert len(logged) == 1 + now = datetime.now() + LiteLLMLoggingObj._get_assembled_streaming_response( + logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[] + ) + assert logged[0].response.usage["prompt_tokens"] == 29 + + client_usage = events[-1].response.usage + assert isinstance(client_usage, ResponseAPIUsage) + assert client_usage.input_tokens == 29 + assert client_usage.cost == pytest.approx(0.0001) From 969d152c4bdc781e706ffffcd496ab1f3332c444 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 16:40:25 -0700 Subject: [PATCH 148/157] fix(jwt): evict jwt_key_mapping cache when a virtual key is deleted The FK cascade drops the LiteLLM_JWTKeyMapping row, but the cached jwt_key_mapping:{claim}:{value} entry still resolved to the deleted token hash, so every JWT call from that identity failed until virtual_key_mapping_cache_ttl expired instead of auto-registering against a recreated key. delete_verification_tokens now snapshots the mapping cache keys before the delete and evicts them across replicas afterwards, the same way /key/regenerate already does. Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- .../key_management_endpoints.py | 19 ++++ .../test_key_management_endpoints.py | 98 +++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f46c4170071..749a940de0e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4559,6 +4559,23 @@ async def delete_verification_tokens( litellm_changed_by=litellm_changed_by, ) + # Snapshot before the delete: the FK cascade drops the mapping rows, but their + # cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380). + jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple( + cache_key + for keys_for_token in await asyncio.gather( + *( + get_jwt_key_mapping_cache_keys_for_token( + hashed_token=key.token, + prisma_client=prisma_client, + ) + for key in authorized_keys + if key.token is not None + ) + ) + for cache_key in keys_for_token + ) + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: deleted_tokens = await prisma_client.delete_data(tokens=tokens) if deleted_tokens is not None and len(deleted_tokens) != len(tokens): @@ -4571,6 +4588,8 @@ async def delete_verification_tokens( if len(deleted_tokens) != len(tokens): failed_tokens = [token for token in tokens if token not in deleted_tokens] + await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + else: raise Exception("DB not connected. prisma_client is None") except Exception as e: 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 d2aeba18f7d..65cc23ea67f 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 @@ -5085,6 +5085,104 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): assert len(deleted_keys) == 2 +class _JWTMappingRow: + def __init__(self, token, jwt_claim_name, jwt_claim_value): + self.token = token + self.jwt_claim_name = jwt_claim_name + self.jwt_claim_value = jwt_claim_value + + +class _CascadingJWTMappingTable: + """Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key is deleted.""" + + def __init__(self, rows): + self.rows = rows + + async def find_many(self, where, **kwargs): + return [row for row in self.rows if row.token == where["token"]] + + def cascade(self, deleted_tokens): + self.rows = [row for row in self.rows if row.token not in deleted_tokens] + + +class _RecordingEvict: + def __init__(self): + self.cache_keys = () + + async def __call__(self, cache_keys, user_api_key_cache): + self.cache_keys = tuple(cache_keys) + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypatch): + """Deleting a key must evict its jwt_key_mapping cache entries (LIT-5380). + + The FK cascade removes the mapping rows, so a surviving cache entry would keep + resolving the deleted token hash and 401 every JWT call from that identity until + virtual_key_mapping_cache_ttl expires, instead of auto-registering again. + """ + jwt_table = _CascadingJWTMappingTable( + [_JWTMappingRow("hashed-token-1", "email", "user@example.com")] + ) + + key1 = LiteLLM_VerificationToken( + token="hashed-token-1", + user_id="user-123", + team_id=None, + key_alias="jwt-mapped-key", + spend=0.0, + max_budget=None, + models=[], + aliases={}, + config={}, + permissions={}, + metadata={}, + model_max_budget={}, + model_spend={}, + soft_budget_cooldown=False, + allowed_routes=[], + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[key1] + ) + mock_prisma_client.db.litellm_jwtkeymapping = jwt_table + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + + async def cascading_delete_data(tokens): + jwt_table.cascade(tokens) + return list(tokens) + + mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data) + + recording_evict = _RecordingEvict() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.evict_and_broadcast", + recording_evict, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", + lambda token: token, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ) + + await delete_verification_tokens( + tokens=["hashed-token-1"], + user_api_key_cache=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + ) + + assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",) + + @pytest.mark.asyncio async def test_delete_key_fn_persists_deleted_keys(monkeypatch): from litellm.proxy._types import KeyRequest From 2000642592670baf38c46f4c6e3bcb851c59d79e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 16:43:50 -0700 Subject: [PATCH 149/157] fix(router): resolve route candidate ids through the router's own resolver (#40491) get_candidate_model_ids_for_route (added in #40280 for the encrypted-content affinity check) reconstructed the candidate pool by unioning the model_name and team indexes with pattern_router.route. That diverged from how the router actually resolves a route: it took a union instead of the first matching path, and pattern_router.route only matches the literal name, so a provider-qualified pattern (matched by get_deployments_by_pattern, which retries the {provider}/{model} form) was missed and the default deployment was ignored. For an affinity follow-up on a wildcard or team-public route, that mismatch could strip encrypted reasoning on a same-group cooldown, or return a 503 on a real cross-path switch. Delegate the non-model_name case to _try_early_resolve_deployments_for_model_not_in_names, the same resolver _common_checks_available_deployment uses, so candidate membership follows the router's real precedence. With include_team_models left off it stays read-only and does not raise. Behavior for concrete model groups and routing groups is unchanged. Claude-Session: https://claude.ai/code/session_01KAumQbhzk6jdWWHFLA8Jar Co-authored-by: Claude Opus 4.8 --- litellm/router.py | 35 +++++++++++++++++-------------- tests/test_litellm/test_router.py | 7 ++++++- 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 88601e5b97b..f9d4bf1428e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11170,28 +11170,31 @@ class Router: def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: """ - Deployment ids that could serve ``model`` for ``team_id``, unioned across the paths - the router resolves a route through: ``model_group_alias``, a routing group, the - ``model_name`` and team indexes, and wildcard pattern routes. Read-only and - side-effect-free, unlike ``_common_checks_available_deployment`` which also applies - fallbacks and can raise. Lets a pre-call check tell a genuine cross-group route from - same-group unavailability without re-deriving that precedence at the call site, and - without leaking deployment ids into request kwargs bound for the provider. + Deployment ids that could serve ``model`` for ``team_id``, following the same + precedence ``_common_checks_available_deployment`` uses to build a candidate pool: + ``model_group_alias``, then a routing group, then the first matching early-resolve + path for a name that is not a ``model_name`` (team route, wildcard pattern via + ``get_deployments_by_pattern``, team pattern router, default deployment), then the + ``model_name`` and team indexes. Delegating to the router's own resolvers keeps this + aligned with how a route actually resolves rather than re-deriving it, and unlike + ``_common_checks_available_deployment`` it is read-only: it does not apply request + fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call + check tell a genuine cross-group route from same-group unavailability without leaking + deployment ids into request kwargs bound for the provider. """ resolved: Final = self._get_model_from_alias(model=model) or model routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) if routing_group_members is not None: return self._deployment_ids(routing_group_members) - if resolved in self.model_names: - return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id)) - team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None - return self._deployment_ids( - ( - *self._get_all_deployments(model_name=resolved, team_id=team_id), - *(self.pattern_router.route(resolved) or ()), - *((team_router.route(resolved) or ()) if team_router is not None else ()), - ) + early: Final = self._try_early_resolve_deployments_for_model_not_in_names( + model=resolved, request_team_id=team_id ) + if early is not None: + early_deployments: Final = early[1] + return self._deployment_ids( + (early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments + ) + return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id)) @staticmethod def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1e0327e5ac2..f5b1f79ad73 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15023,7 +15023,9 @@ def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): pre-call check can tell a genuine cross-group route from same-group unavailability. A concrete model group returns its member ids; a wildcard/pattern deployment is included for a concrete model it matches, which the bare model_name index misses. - Regression guard for the LIT-7195 tier-change discriminator's team/pattern gaps. + The unprefixed-name case must resolve through get_deployments_by_pattern (which retries + the provider-qualified form), not a bare pattern_router.route that only sees the literal + name. Regression guard for the LIT-7195 tier-change discriminator's team/pattern gaps. """ router = Router( model_list=[ @@ -15047,6 +15049,9 @@ def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): assert router.get_candidate_model_ids_for_route(model="grp") == frozenset({"dep-a", "dep-b"}) assert "dep-wild" in router.get_candidate_model_ids_for_route(model="openai/gpt-4o-some-new-model") + # unprefixed name whose provider resolves to openai: only get_deployments_by_pattern's + # provider-qualified retry matches "openai/*"; a bare route() on the literal name misses it + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="gpt-5") def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): From 70c238dde0c4bef5cbb976fbd382320a8021352d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:54:43 -0700 Subject: [PATCH 150/157] fix(ui): preserve dotted MCP tool argument names --- ...MCPToolArgumentsForm.integration.test.tsx} | 114 +++++++++++++++++- .../mcp_tools/MCPToolArgumentsForm.tsx | 44 ++++--- 2 files changed, 141 insertions(+), 17 deletions(-) rename ui/litellm-dashboard/src/components/mcp_tools/{MCPToolArgumentsForm.test.tsx => MCPToolArgumentsForm.integration.test.tsx} (61%) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx similarity index 61% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx rename to ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx index d6e614da5eb..a4302f9fd6a 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { render, screen } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect } from "vitest"; import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "./MCPToolArgumentsForm"; @@ -26,6 +26,118 @@ const submitError = async (ref: React.RefObject) }; describe("MCPToolArgumentsForm", () => { + it("keeps dotted arguments separate from a same-prefix object and converts their values", async () => { + const ref = renderForm({ + type: "object", + properties: { + "filter.category": { type: "string" }, + filter: { type: "object" }, + "page.limit": { type: "integer" }, + query: { type: "string" }, + }, + required: ["filter.category"], + }); + + fireEvent.change(screen.getByRole("textbox", { name: "filter.category *" }), { + target: { value: "invoices" }, + }); + fireEvent.change(screen.getByRole("textbox", { name: "filter", exact: true }), { + target: { value: '{"category":"receipts","metadata":{"region":"eu"}}' }, + }); + fireEvent.change(screen.getByRole("spinbutton", { name: "page.limit" }), { target: { value: "7" } }); + fireEvent.change(screen.getByRole("textbox", { name: "query" }), { target: { value: "September" } }); + + const expected = { + "filter.category": "invoices", + filter: { category: "receipts", metadata: { region: "eu" } }, + "page.limit": 7, + query: "September", + }; + await expect(submit(ref)).resolves.toEqual(expected); + }); + + it("shows required validation on the literal dotted field and accepts a correction", async () => { + const ref = renderForm({ + type: "object", + properties: { "filter.category": { type: "string" } }, + required: ["filter.category"], + }); + + expect(await submitError(ref)).toEqual({ + errorFields: [{ name: ["filter.category"], errors: ["Please enter filter.category"] }], + }); + expect(await screen.findByText("Please enter filter.category")).toBeInTheDocument(); + expect(screen.getByRole("textbox", { name: "filter.category *" })).toHaveAttribute("aria-invalid", "true"); + + fireEvent.change(screen.getByRole("textbox", { name: "filter.category *" }), { + target: { value: "invoices" }, + }); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "invoices" }); + }); + + it("validates JSON for dotted arguments inside params and preserves their literal names", async () => { + const ref = renderForm({ + type: "object", + properties: { + params: { + type: "object", + properties: { "filter.options": { type: "object" } }, + required: ["filter.options"], + }, + }, + required: [], + }); + const field = screen.getByRole("textbox", { name: "filter.options *" }); + fireEvent.change(field, { target: { value: "invalid" } }); + + expect(await submitError(ref)).toEqual({ + errorFields: [{ name: ["filter.options"], errors: ["Invalid JSON"] }], + }); + expect(await screen.findByText("Invalid JSON")).toBeInTheDocument(); + + fireEvent.change(field, { target: { value: '{"region":"eu"}' } }); + await expect(submit(ref)).resolves.toEqual({ params: { "filter.options": { region: "eu" } } }); + }); + + it("resets dotted defaults and positional values when the selected tool changes", async () => { + const ref = React.createRef(); + const { rerender } = render( + , + ); + expect(screen.getByRole("textbox", { name: "filter.category" })).toHaveValue("invoices"); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "invoices" }); + fireEvent.change(screen.getByRole("textbox", { name: "filter.category" }), { + target: { value: "edited" }, + }); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "edited" }); + + rerender( + , + ); + expect(screen.getByRole("textbox", { name: "filter.category" })).toHaveValue("receipts"); + await expect(submit(ref)).resolves.toEqual({ query: "new tool", "filter.category": "receipts" }); + }); + it("returns typed values for a string, integer, number and boolean field", async () => { const user = userEvent.setup(); const ref = renderForm({ diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx index ca8c5697e6e..4fff3e57fd9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx @@ -9,7 +9,10 @@ import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { MCPTool, InputSchema, InputSchemaProperty } from "./types"; -type ToolFormValues = Record; +type ToolFormValues = { args: unknown[] }; + +const argumentValues = (schema: InputSchema, values: ToolFormValues): Record => + Object.fromEntries(Object.keys(schema.properties ?? {}).map((key, index) => [key, values.args[index]])); const STRING_SCHEMA_MESSAGES: Readonly> = { input: "Please enter input for this tool" }; @@ -38,7 +41,7 @@ type FieldError = { type: string; message: string }; const collectErrors = ( actualSchema: InputSchema, requiredMessages: Readonly>, - values: ToolFormValues, + values: Record, ): Record => { const entries = Object.entries(actualSchema.properties ?? {}).flatMap<[string, FieldError]>(([key, prop]) => { const value = values[key]; @@ -57,8 +60,18 @@ const collectErrors = ( const buildResolver = (actualSchema: InputSchema, requiredMessages: Readonly> = {}): Resolver => (values) => { - const errors = collectErrors(actualSchema, requiredMessages, values); - return Object.keys(errors).length > 0 ? { values: {}, errors } : { values, errors: {} }; + const errors = collectErrors(actualSchema, requiredMessages, argumentValues(actualSchema, values)); + if (Object.keys(errors).length === 0) return { values, errors: {} }; + return { + values: {}, + errors: { + args: Object.fromEntries( + Object.keys(actualSchema.properties ?? {}).flatMap((key, index) => + Object.hasOwn(errors, key) ? [[index, errors[key]]] : [], + ), + ), + }, + }; }; const labelFor = (key: string, prop: InputSchemaProperty, required: boolean): React.ReactNode => ( @@ -238,10 +251,7 @@ const MCPToolArgumentsForm = forwardRef( - () => - Object.fromEntries( - Object.entries(actualSchema.properties ?? {}).map(([key, prop]) => [key, getInitialValueForField(prop)]), - ), + () => ({ args: Object.values(actualSchema.properties ?? {}).map(getInitialValueForField) }), [actualSchema], ); @@ -255,7 +265,7 @@ const MCPToolArgumentsForm = forwardRef ({ getSubmitValues: async () => { - const values = form.getValues(); + const values = argumentValues(actualSchema, form.getValues()); const errors = collectErrors(actualSchema, requiredMessages, values); if (Object.keys(errors).length > 0) { await form.trigger(); @@ -286,14 +296,16 @@ const MCPToolArgumentsForm = forwardRef Input * } > - {(field) => } + {(field) => ( + + )} @@ -318,13 +330,13 @@ const MCPToolArgumentsForm = forwardRef - {Object.entries(actualSchema.properties).map(([key, prop]) => { + {Object.entries(actualSchema.properties).map(([key, prop], index) => { const required = actualSchema.required?.includes(key) ?? false; return ( {(field) => { @@ -375,7 +387,7 @@ const MCPToolArgumentsForm = forwardRef ); @@ -385,7 +397,7 @@ const MCPToolArgumentsForm = forwardRef ); From 0b21b99ebd193872c520996cec79defad07ee108 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:57:15 -0700 Subject: [PATCH 151/157] fix(ui): type MCP argument resolver results explicitly --- .../mcp_tools/MCPToolArgumentsForm.integration.test.tsx | 2 +- .../src/components/mcp_tools/MCPToolArgumentsForm.tsx | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx index a4302f9fd6a..b1174d1d37d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx @@ -41,7 +41,7 @@ describe("MCPToolArgumentsForm", () => { fireEvent.change(screen.getByRole("textbox", { name: "filter.category *" }), { target: { value: "invoices" }, }); - fireEvent.change(screen.getByRole("textbox", { name: "filter", exact: true }), { + fireEvent.change(screen.getByRole("textbox", { name: "filter" }), { target: { value: '{"category":"receipts","metadata":{"region":"eu"}}' }, }); fireEvent.change(screen.getByRole("spinbutton", { name: "page.limit" }), { target: { value: "7" } }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx index 4fff3e57fd9..ab3213d9b37 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx @@ -1,6 +1,6 @@ import React, { forwardRef, useImperativeHandle, useMemo } from "react"; import { CircleHelp } from "lucide-react"; -import { useForm, type Resolver } from "react-hook-form"; +import { useForm, type Resolver, type ResolverResult } from "react-hook-form"; import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; @@ -59,7 +59,7 @@ const collectErrors = ( const buildResolver = (actualSchema: InputSchema, requiredMessages: Readonly> = {}): Resolver => - (values) => { + (values): ResolverResult => { const errors = collectErrors(actualSchema, requiredMessages, argumentValues(actualSchema, values)); if (Object.keys(errors).length === 0) return { values, errors: {} }; return { From 6b721de3e530bc97758d59621e57e3a75a9ebadf Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 17:15:48 -0700 Subject: [PATCH 152/157] fix(databricks): route Unity model services through AI Gateway (#40492) Co-authored-by: Claude Code --- .../llms/databricks/chat/transformation.py | 6 +- litellm/llms/databricks/common_utils.py | 20 ++++-- .../test_databricks_chat_transformation.py | 13 ++++ .../test_databricks_partner_integration.py | 71 +++++++++++++++++++ 4 files changed, 101 insertions(+), 9 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 54c7040daf9..09aaf970dc5 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -250,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = self._get_api_base(api_base) - complete_url: Final = f"{api_base}/chat/completions" + use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2 + api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) + url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base + complete_url: Final = f"{url_base}/chat/completions" return complete_url def get_supported_openai_params(self, model: str | None = None) -> list: diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 7695b1cb35e..a4ec2c5378b 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -177,19 +177,13 @@ class DatabricksBase: # Default: just litellm return f"litellm/{version}" - def _get_api_base(self, api_base: str | None) -> str: - """ - Get the Databricks API base URL. - - If not provided, attempts to get it from the Databricks SDK. - """ + def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str: if api_base is None: try: from databricks.sdk import WorkspaceClient databricks_client: Final = WorkspaceClient() api_base = f"{databricks_client.config.host}/serving-endpoints" - return api_base except ImportError: raise DatabricksException( status_code=400, @@ -198,6 +192,18 @@ class DatabricksBase: "or install the databricks-sdk Python library." ), ) + + if not use_ai_gateway: + return api_base + + normalized_api_base: Final = api_base.rstrip("/") + if normalized_api_base.endswith("/ai-gateway/mlflow/v1"): + return normalized_api_base + if normalized_api_base.endswith("/serving-endpoints"): + return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1" + api_base_parts: Final = urlsplit(normalized_api_base) + if api_base_parts.path in ("", "/"): + return f"{normalized_api_base}/ai-gateway/mlflow/v1" return api_base def _get_oauth_m2m_token( diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 38d5844b9b9..caf7bed7385 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -255,6 +255,19 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_request_preserves_unity_model_service_name(): + config = DatabricksConfig() + result = config.transform_request( + model="system.ai.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "system.ai.kimi-k3" + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 39198bb20f3..c6ad78366f7 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -657,6 +657,77 @@ class TestEndpointURLConstruction: assert api_base.endswith("/chat/completions") + def test_chat_gateway_endpoint_for_unity_model_on_legacy_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_endpoint_preserves_explicit_gateway_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1/", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_preserves_unity_model_service_name_with_explicit_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + config = DatabricksConfig() + request = config.transform_request( + model="catalog.schema.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert config.get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1", + api_key="test-key", + model="catalog.schema.kimi-k3", + optional_params={}, + litellm_params={}, + ) == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + assert request["model"] == "catalog.schema.kimi-k3" + + def test_chat_legacy_endpoint_remains_default(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="databricks-kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/serving-endpoints/chat/completions" + def test_embeddings_endpoint(self, monkeypatch): """Embeddings endpoint is correctly appended.""" monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) From b8d7f68aeba2bd4a599895c910df085c2eb530a1 Mon Sep 17 00:00:00 2001 From: kerry-berri Date: Wed, 9 Sep 2026 17:46:12 -0700 Subject: [PATCH 153/157] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../src/components/llm_calls/chat_completion.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index b813a35f687..c12a2c3cfb4 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -245,7 +245,7 @@ export async function makeOpenAIChatCompletionRequest( // Extract cost from usage object if available if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - const parsedCost = parseFloat(chunkWithUsage.usage.cost); + const parsedCost = Number(chunkWithUsage.usage.cost); if (Number.isFinite(parsedCost)) { usageData.cost = parsedCost; } From 317b29e69d658ca2502e656d4a30c4db5ced0198 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 18:02:23 -0700 Subject: [PATCH 154/157] feat(cli): add lite configure claude and lite unconfigure claude (#40319) Persistently route Claude Code through a LiteLLM proxy with a long-lived virtual key or the stored lite login, turn on gateway model discovery so /model lists the proxy's models, optionally pick the model Claude Code starts on, and record what changed so unconfigure restores only the keys the user has not touched since. lite login --config-claude writes through the same receipt and is undoable too. The two settings merges (lite up / --config-claude and lite autoroute) collapse into one credential-aware merge --- litellm/proxy/client/cli/README.md | 22 +- litellm/proxy/client/cli/commands/agents.py | 5 +- litellm/proxy/client/cli/commands/auth.py | 29 +- .../client/cli/commands/autoroute/commands.py | 9 +- .../client/cli/commands/autoroute/settings.py | 51 -- .../client/cli/commands/claude_settings.py | 517 +++++++++++++++-- .../proxy/client/cli/commands/configure.py | 252 +++++++++ litellm/proxy/client/cli/commands/pi.py | 31 +- litellm/proxy/client/cli/commands/up.py | 9 +- litellm/proxy/client/cli/main.py | 4 + .../client/cli/autoroute/test_commands.py | 4 + .../client/cli/autoroute/test_settings.py | 63 --- .../proxy/client/cli/test_auth_commands.py | 40 +- .../proxy/client/cli/test_claude_settings.py | 535 +++++++++++++++++- .../client/cli/test_configure_commands.py | 331 +++++++++++ .../test_litellm/proxy/client/cli/test_pi.py | 26 +- .../proxy/client/cli/test_up_commands.py | 40 +- 17 files changed, 1718 insertions(+), 250 deletions(-) delete mode 100644 litellm/proxy/client/cli/commands/autoroute/settings.py create mode 100644 litellm/proxy/client/cli/commands/configure.py delete mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_settings.py create mode 100644 tests/test_litellm/proxy/client/cli/test_configure_commands.py diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index f7d9eb7da9a..03d01ff7f66 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. @@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops. +#### Configuring Claude Code Once, With a Virtual Key or Your Login + +`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh +lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto +claude +``` + +With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (the ones whose id contains `claude` or `anthropic`) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control + +Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt + +What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request + ### QA Complexity-Based Auto-Routing Against Your Real Proxy `lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. @@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. -You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) +You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 7a0ae9dc955..b287293eb7a 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -16,6 +16,7 @@ from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, + ListingFailure, PiSyncError, fetch_model_ids, fetch_model_limits, @@ -165,7 +166,9 @@ def prepare_pi( """ ids: Final = fetch_model_ids(base_url, api_key, get=get) if isinstance(ids, PiSyncError): - raise AgentRunError(ids.message) + raise AgentRunError( + f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message + ) limits: Final = fetch_model_limits(base_url, api_key, get=get) path: Final = models_json_path(base_env) error: Final = sync_models_json(path, base_url, ids, limits) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 12a288202b6..24b1c766664 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -41,9 +41,15 @@ from litellm.litellm_core_utils.cli_token_utils import ( from .claude_settings import ( CLAUDE_SETTINGS_PATH, + CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS, + STARTING_MODEL_ROLE, + ApiKeyHelper, ClaudeSettingsError, - write_claude_settings, + KeepModel, + configure_claude_settings, + refuse_while_owned, + resolve_api_key_helper, ) from .pkce_login import ( Http, @@ -778,13 +784,23 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: def _configure_claude_code(base_url: str) -> None: - """Point Claude Code at base_url by patching ~/.claude/settings.json.""" + """Point Claude Code at base_url by patching ~/.claude/settings.json, undoable with `lite unconfigure claude`.""" try: - write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + configure_claude_settings( + base_url, + ApiKeyHelper(resolve_api_key_helper(base_url)), + KeepModel(), + CLAUDE_SETTINGS_PATH, + CONFIGURE_STATE_PATH, + SETTINGS_FILE_OWNERS, + ) except ClaudeSettingsError as e: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.") - click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.") + click.echo( + "Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. " + f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}." + ) def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None: @@ -853,6 +869,11 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] + if config_claude: + try: + refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + except ClaudeSettingsError as e: + raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}") try: if pkce: diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 26d45138a27..86701f186fd 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -14,11 +14,13 @@ from ..claude_settings import ( AUTOROUTE_BACKUP_PATH, CLAUDE_SETTINGS_PATH, ClaudeSettingsError, + StaticToken, load_json_or_empty, + merge_claude_settings, ) from ..up import BackupRecord as ClaudeBackupRecord from ..up import restore_claude_settings, write_backup -from .config import master_key_from_config +from .config import AUTOROUTER_MODEL_NAME, master_key_from_config from .process import ( CONFIG_PATH, DEFAULT_AUTOROUTE_PORT, @@ -37,7 +39,6 @@ from .process import ( terminate, write_pid_record, ) -from .settings import merge_claude_settings_static_token from .wizard import run_configure_wizard _GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) @@ -156,7 +157,9 @@ def up(port: int) -> None: ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), AUTOROUTE_BACKUP_PATH, ) - merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key) + merged: Final = merge_claude_settings( + original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME + ) CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) with secure_create(CLAUDE_SETTINGS_PATH) as f: json.dump(merged, f, indent=2) diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py deleted file mode 100644 index 60729b5410d..00000000000 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import Final - -from pydantic import JsonValue - -from .config import AUTOROUTER_MODEL_NAME - -ENV_KEY: Final = "env" -API_KEY_HELPER_KEY: Final = "apiKeyHelper" -ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" -ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" -ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" -ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" -ENABLE_TOOL_SEARCH_VALUE: Final = "true" -# Force every one of Claude Code's own model tiers to request the auto-router by name. -# Router's auto-router registry is keyed by the literal requested model string -# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" -# model_name can never work as a catch-all -- these overrides are what actually makes -# Claude Code send "autorouter" regardless of /model or its own version-specific defaults. -ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( - "ANTHROPIC_DEFAULT_SONNET_MODEL", - "ANTHROPIC_DEFAULT_HAIKU_MODEL", - "ANTHROPIC_DEFAULT_OPUS_MODEL", -) - - -def merge_claude_settings_static_token( - settings: dict[str, JsonValue], base_url: str, auth_token: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to a local ephemeral proxy with a static token. - - Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real - remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the - locally persisted autoroute master key, so a plain env var is simpler and correct. Any - existing apiKeyHelper is cleared so it can't fight with the static token. - """ - raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final[dict[str, JsonValue]] = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - **base_env, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - ANTHROPIC_AUTH_TOKEN_KEY: auth_token, - **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, - } - env.pop(ANTHROPIC_API_KEY_KEY, None) - merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env} - merged.pop(API_KEY_HELPER_KEY, None) - return merged - - -__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 5e3ce95f088..bd9cf85dc59 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -1,37 +1,70 @@ """Shared handling of Claude Code's ~/.claude/settings.json. -`lite up` patches this file temporarily and restores it on exit; `lite login ---config-claude` patches it persistently. Both need the same merge and the same -apiKeyHelper command, and `up` already imports from `auth`, so the shared parts -live here rather than in either command module. +`lite up` and `lite autoroute up` patch this file temporarily and restore it on +exit; `lite login --config-claude` and `lite configure claude` patch it +persistently and record how to undo it. All of them need the same merge and the +same apiKeyHelper command, and `up` already imports from `auth`, so the shared +parts live here rather than in any one command module. """ +import hashlib +import json import shlex import shutil import sys -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import reduce +from itertools import chain from pathlib import Path -from typing import Final +from types import MappingProxyType +from typing import Final, TypeAlias -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError -from litellm.litellm_core_utils.private_json import write_private_json +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_json, +) from .cmd_quoting import quote_for_cmd ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" +MODEL_KEY: Final = "model" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" +ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", +) +OWNED_ENV_KEYS: Final = ( + ENABLE_TOOL_SEARCH_KEY, + ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, + ANTHROPIC_BASE_URL_KEY, + ANTHROPIC_AUTH_TOKEN_KEY, + ANTHROPIC_API_KEY_KEY, +) +OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY) +OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS) +_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY)) +_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY) +_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}" +STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" +CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" @dataclass(frozen=True, slots=True) @@ -55,6 +88,102 @@ class ClaudeSettingsError(Exception): """Raised for any user-actionable failure while reading or writing Claude Code settings.""" +@dataclass(frozen=True, slots=True) +class StaticToken: + """A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN.""" + + token: str + + +@dataclass(frozen=True, slots=True) +class ApiKeyHelper: + """A `lite auth print-token` command Claude Code runs per request, so a login renews in place.""" + + command: str + + +ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper + + +@dataclass(frozen=True, slots=True) +class KeepModel: + """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login).""" + + +@dataclass(frozen=True, slots=True) +class UnpinModel: + """Let go of a `model` an earlier configure pinned; one the user set themselves stays.""" + + +@dataclass(frozen=True, slots=True) +class StartOn: + """Pin the top-level `model`, the row Claude Code starts on.""" + + model: str + + +ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn + + +class OwnedValue(BaseModel): + """What one key held at a moment in time; `present=False` is an absent key, not a null one.""" + + model_config = ConfigDict(frozen=True) + + present: bool + value: JsonValue = None + + +class ConfigureReceipt(BaseModel): + """What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key). + + Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the + value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and + carries the earlier one otherwise, so a key the user edited in between stops matching and is left + alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier + snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the + repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot + was captured beside, so a credential is only ever put back next to the server it was issued for. + No fingerprint is a second copy of a token. + """ + + model_config = ConfigDict(frozen=True) + + file_existed: bool + env_present: bool + env_was_object: bool + previous: Mapping[str, OwnedValue] + written: Mapping[str, str] + endpoints: Mapping[str, OwnedValue] + + +@dataclass(frozen=True, slots=True) +class WithheldCredential: + """A credential left removed: captured beside `endpoint`, while the restored file points elsewhere.""" + + key: str + endpoint: str + + +@dataclass(frozen=True, slots=True) +class UnconfigureOutcome: + """Keys whose value unconfigure changed back, keys the user changed since and so were left as they + are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the + URL points back), and whether no settings file remains.""" + + restored: tuple[str, ...] + kept: tuple[str, ...] + withheld: tuple[WithheldCredential, ...] = () + file_removed: bool = False + + +@dataclass(frozen=True, slots=True) +class _Claim: + previous: OwnedValue + written: str | None + endpoint: OwnedValue | None + + def load_json_or_empty(path: Path) -> dict[str, JsonValue]: try: content: Final = path.read_bytes() if path.exists() else b"" @@ -70,29 +199,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]: ) -def merge_claude_settings( - settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to route Claude Code through the proxy. +def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + if raw_env is None: + return MappingProxyType({}) + if not isinstance(raw_env, dict): + raise ClaudeSettingsError( + f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.' + ) + return raw_env - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH - defaults to true because Claude Code turns tool search off when - ANTHROPIC_BASE_URL is not a first-party Anthropic host, and - CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker - is filled from the proxy's /v1/models; existing values of both are left - alone. Every other key is preserved untouched. + +def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: + """Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a + purely local check, so commands run it before any login prompt or request.""" + for owner in owners: + if owner.backup_path.exists(): + raise ClaudeSettingsError( + f"`{owner.start_command}` is currently managing {settings_path} (backup at " + f"{owner.backup_path}) and will restore it when it stops. " + f"Run `{owner.stop_command}` first, then retry." + ) + + +def _write_target(settings_path: Path) -> Path: + """Write through a symlinked settings.json rather than replacing the link, which would silently + detach a file symlinked into a dotfiles repo.""" + try: + return settings_path.resolve() if settings_path.is_symlink() else settings_path + except OSError as e: + raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e + + +def _stage(path: Path, document: Mapping[str, object]) -> str: + try: + return stage_private_json(str(path), document) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {path}: {e}") from e + + +def _land( + path: Path, + staged: str | None, + also_discard: Sequence[str | None] = (), + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a + filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are + discarded, so no temp file holding a token is left behind.""" + try: + if staged is None: + path.unlink(missing_ok=True) + else: + commit(staged, str(path)) + except OSError as e: + for other in also_discard: + if other is not None: + discard_staged_json(other) + raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], + base_url: str, + credential: ClaudeCredential, + default_model: str | None = None, + tier_model: str | None = None, +) -> Mapping[str, JsonValue]: + """Return a new settings mapping wired to route Claude Code through the proxy. + + A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper; + the other credential slots are removed either way, since Claude Code given two credentials may + send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their + defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts + on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one + group. Apart from those tier keys, exactly OWNED_PATHS are touched. """ raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE, - **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - } - return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} + current_env: Final = raw_env if isinstance(raw_env, dict) else {} + env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ( + (ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE), + (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE), + ), + ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS), + ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),), + ((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (), + ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None), + ) + ) + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)), + ((ENV_KEY, env),), + ((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (), + ((MODEL_KEY, default_model),) if default_model is not None else (), + ) + ) def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: @@ -121,56 +325,255 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) -def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: - """Persistently point Claude Code at base_url, preserving every unrelated setting. +def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue: + return OwnedValue(present=key in container, value=container.get(key)) - Refuses while any owner holds a backup: each restores its backup when it - stops, which would silently undo this write. - """ - for owner in owners: - if owner.backup_path.exists(): - raise ClaudeSettingsError( - f"`{owner.start_command}` is currently managing {settings_path} (backup at " - f"{owner.backup_path}) and will restore it when it stops. " - f"Run `{owner.stop_command}` first, then retry." - ) - normalized_base_url: Final = base_url.rstrip("/") - api_key_helper: Final = resolve_api_key_helper(normalized_base_url) - existing: Final = load_json_or_empty(settings_path) - raw_env: Final = existing.get(ENV_KEY) - if raw_env is not None and not isinstance(raw_env, dict): - raise ClaudeSettingsError( - f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. ' - "Fix or remove it, then retry." - ) - merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper) - # os.replace() swaps the symlink itself for a regular file, silently detaching a - # settings.json that is symlinked into a dotfiles repo. There is no backup to undo - # that here, unlike `lite up`, so write through to the link's target instead. - target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path + +def _fingerprint(owned: OwnedValue) -> str: + return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest() + + +def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + return raw_env if isinstance(raw_env, dict) else MappingProxyType({}) + + +def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue: + section, _, key = path.rpartition(".") + return _owned(_env(settings) if section else settings, key) + + +def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ()) + ) + + +def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + """`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes.""" + section, _, key = path.rpartition(".") + if not section: + return _with_key(settings, key, owned) + return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned))) + + +def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]: + return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings) + + +def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool: + """Whether the key still holds what a configure wrote (a key no configure ever changed is never ours).""" + return receipt.written.get(path) == _fingerprint(_lookup(settings, path)) + + +def _claim( + path: str, + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + url_now: OwnedValue, +) -> _Claim: + """What this configure records for one key; see ConfigureReceipt for the rules.""" + before, after = _lookup(current, path), _lookup(merged, path) + carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None + return _Claim( + previous=before if carried is None else carried.previous.get(path, before), + written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)), + endpoint=None + if path not in _CREDENTIAL_PATHS + else (url_now if carried is None else carried.endpoints.get(path, url_now)), + ) + + +def _receipt( + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + file_exists: bool, +) -> ConfigureReceipt: + url_now: Final = _lookup(current, _BASE_URL_PATH) + claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS}) + return ConfigureReceipt( + file_existed=file_exists if earlier is None else earlier.file_existed, + env_present=ENV_KEY in current if earlier is None else earlier.env_present, + env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object, + previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}), + written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}), + endpoints=MappingProxyType( + {path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None} + ), + ) + + +def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: + if not state_path.exists(): + return None try: - write_private_json(str(target), merged) + return ConfigureReceipt.model_validate_json(state_path.read_bytes()) + except (OSError, ValidationError) as e: + raise ClaudeSettingsError( + f"{state_path} is not a readable `lite configure claude` receipt ({e}). " + "Remove it and edit Claude Code's settings by hand if they still point at the proxy." + ) from e + + +def configure_claude_settings( + base_url: str, + credential: ClaudeCredential, + model: ModelChoice, + settings_path: Path, + state_path: Path, + owners: Sequence[SettingsFileOwner], + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Persistently route Claude Code through base_url, recording how to undo it. + + Both files are staged before either is committed, so a full disk or a read-only directory fails + before anything changes. The two commits are still two renames: a receipt rename that fails + discards the staged settings, and a settings rename that fails after the receipt landed puts the + earlier receipt back (or removes the new one), so the receipt on disk never describes settings + that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an + earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login). + """ + refuse_while_owned(settings_path, owners) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + earlier: Final = read_configure_receipt(state_path) + existing: Final = ( + _with(current, MODEL_KEY, earlier.previous[MODEL_KEY]) + if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier) + else current + ) + merged: Final = merge_claude_settings( + existing, base_url, credential, model.model if isinstance(model, StartOn) else None + ) + receipt: Final = _receipt(current, merged, earlier, settings_path.exists()) + target: Final = _write_target(settings_path) + try: + ensure_private_dir(state_path.parent) except OSError as e: - raise ClaudeSettingsError(f"Could not write {target}: {e}") from e + raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e + staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json")) + try: + staged_settings: Final = _stage(target, merged) + except ClaudeSettingsError: + discard_staged_json(staged_receipt) + raise + _land(state_path, staged_receipt, (staged_settings,), commit) + try: + _land(target, staged_settings, commit=commit) + except ClaudeSettingsError as settings_error: + try: + _land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json"))) + except ClaudeSettingsError as receipt_error: + raise ClaudeSettingsError( + f"{settings_error} The receipt at {state_path} now describes settings that were not written and " + f"could not be put back either ({receipt_error}); remove it before retrying." + ) from settings_error + raise + + +def _endpoint_text(endpoint: OwnedValue) -> str: + if not endpoint.present: + return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)" + return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value) + + +def unconfigure_claude_settings( + settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner] +) -> UnconfigureOutcome: + """Undo `lite configure claude`: put back every key still holding what configure wrote, leave the + rest alone, and withhold a credential the restored file would send to a different server than it + was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish).""" + refuse_while_owned(settings_path, owners) + receipt: Final = read_configure_receipt(state_path) + if receipt is None: + raise ClaudeSettingsError( + f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo." + ) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt)) + kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present) + put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours})) + url_after: Final = _lookup(put_back, _BASE_URL_PATH) + withheld: Final = tuple( + WithheldCredential(path, _endpoint_text(receipt.endpoints[path])) + for path in _CREDENTIAL_PATHS + if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after + ) + absent: Final = OwnedValue(present=False) + trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld})) + settings: Final = ( + trimmed + if _env(trimmed) or receipt.env_was_object + else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None)) + ) + target: Final = _write_target(settings_path) + file_removed: Final = not settings and not (receipt.file_existed and target.exists()) + kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy + receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}}) + if withheld + else None + ) + staged_settings: Final = None if file_removed else _stage(target, settings) + try: + staged_receipt: Final = ( + None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json")) + ) + except ClaudeSettingsError: + if staged_settings is not None: + discard_staged_json(staged_settings) + raise + _land(target, staged_settings, (staged_receipt,)) + _land(state_path, staged_receipt) + return UnconfigureOutcome( + restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)), + kept=kept, + withheld=withheld, + file_removed=file_removed, + ) __all__ = ( "ANTHROPIC_API_KEY_KEY", + "ANTHROPIC_AUTH_TOKEN_KEY", "ANTHROPIC_BASE_URL_KEY", + "ANTHROPIC_DEFAULT_MODEL_ENV_KEYS", "API_KEY_HELPER_KEY", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", + "CONFIGURE_STATE_PATH", "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY", "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE", "ENABLE_TOOL_SEARCH_KEY", "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", + "MODEL_KEY", + "OWNED_ENV_KEYS", + "OWNED_PATHS", + "OWNED_TOP_LEVEL_KEYS", "SETTINGS_FILE_OWNERS", + "STARTING_MODEL_ROLE", + "ApiKeyHelper", + "ClaudeCredential", "ClaudeSettingsError", + "ConfigureReceipt", + "KeepModel", + "ModelChoice", + "OwnedValue", "SettingsFileOwner", + "StartOn", + "StaticToken", + "UnconfigureOutcome", + "UnpinModel", + "WithheldCredential", + "configure_claude_settings", "load_json_or_empty", "merge_claude_settings", + "read_configure_receipt", + "refuse_while_owned", "resolve_api_key_helper", - "write_claude_settings", + "unconfigure_claude_settings", ) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py new file mode 100644 index 00000000000..f3a98a48476 --- /dev/null +++ b/litellm/proxy/client/cli/commands/configure.py @@ -0,0 +1,252 @@ +"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" + +import re +import sys +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Final + +import click +from InquirerPy import inquirer +from InquirerPy.base.control import Choice + +from .auth import CliContextObj, context_secret_vault, get_stored_api_key +from .claude_settings import ( + CLAUDE_SETTINGS_PATH, + CONFIGURE_STATE_PATH, + SETTINGS_FILE_OWNERS, + STARTING_MODEL_ROLE, + ApiKeyHelper, + ClaudeCredential, + ClaudeSettingsError, + ModelChoice, + StartOn, + StaticToken, + UnconfigureOutcome, + UnpinModel, + configure_claude_settings, + refuse_while_owned, + resolve_api_key_helper, + unconfigure_claude_settings, +) +from .pi import ListingFailure, PiSyncError, fetch_model_ids +from .up import ensure_fresh_login + +_LISTED_MODELS_SHOWN: Final = 20 +_CLAUDE_TARGET: Final = "claude" +_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) +_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" +_CLAUDE_CODE_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +_MODEL_OPTION_HELP: Final = ( + f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, " + "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude " + "Code's sub-agent or background tiers; `lite autoroute up` is the mode that does." +) + + +def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]: + """The credential to write and the key to check the proxy with. + + An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes + into settings.json as a static token. Without one, the stored `lite login` credential is used + the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a + day and renews in place there; a missing or stale login is refreshed first, as `lite up` does. + """ + ctx_obj: Final[CliContextObj] = ctx.obj + explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) + if explicit: + return StaticToken(explicit), explicit + base_url: Final = ctx_obj["base_url"] + ensure_fresh_login(ctx) + stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) + if not stored: + raise ClaudeSettingsError("Login did not produce a usable token.") + return ApiKeyHelper(resolve_api_key_helper(base_url)), stored + + +def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]: + """Every configure path begins the same way: the local ownership check first, so a `lite up` + session is refused before any login prompt or request, then the credential, then the listing.""" + try: + refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + credential, key = resolve_credential(ctx, api_key) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + return credential, _listed_models(ctx.obj["base_url"], key) + + +def _listing_error(base_url: str, error: PiSyncError) -> str: + """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" + if error.kind is ListingFailure.REJECTED: + return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key." + if error.kind is ListingFailure.UNREACHABLE: + return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" + if error.kind is ListingFailure.EMPTY: + return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model." + return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." + + +def _listed_models(base_url: str, key: str) -> tuple[str, ...]: + listed: Final = fetch_model_ids(base_url, key) + if isinstance(listed, PiSyncError): + raise click.ClickException(_listing_error(base_url, listed)) + return listed + + +def _model_choice(model: str | None) -> ModelChoice: + return StartOn(model) if model is not None else UnpinModel() + + +def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequence[str], model: str | None) -> None: + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] + if model is not None and model not in listed: + shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) + more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" + raise click.ClickException( + f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}." + ) + try: + configure_claude_settings( + base_url, credential, _model_choice(model), CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS + ) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + in_picker: Final = sum(1 for listed_model in listed if _CLAUDE_CODE_PICKER_FILTER.search(listed_model)) + click.echo(f"Configured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url}.") + click.echo( + "Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN." + if isinstance(credential, StaticToken) + else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it." + ) + click.echo( + f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model." + if model is not None + else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " + "pass --model to start on a proxy model." + ) + click.echo( + f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing " + "'claude' or 'anthropic')." + ) + click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") + if isinstance(credential, StaticToken) and CLAUDE_SETTINGS_PATH.is_symlink(): + click.echo( + f"Note: {CLAUDE_SETTINGS_PATH} is a symlink to {CLAUDE_SETTINGS_PATH.resolve()}, so your key now lives in " + "that file; keep it out of version control.", + err=True, + ) + + +def _pick_targets() -> tuple[str, ...]: + picked: Final = inquirer.checkbox( + message="Which agents should route through LiteLLM?", + choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS], + validate=lambda chosen: len(chosen) > 0, + invalid_message="Pick at least one.", + ).execute() + return tuple(str(value) for value in picked) + + +def _pick_model(listed: Sequence[str]) -> str | None: + picked: Final = inquirer.fuzzy( + message="Model Claude Code starts on (type to filter; /model switches any time):", + choices=[_KEEP_DEFAULT_MODEL, *listed], + ).execute() + return None if picked == _KEEP_DEFAULT_MODEL else str(picked) + + +def interactive_configure( + ctx: click.Context, + pick_targets: Callable[[], tuple[str, ...]] = _pick_targets, + pick_model: Callable[[Sequence[str]], str | None] = _pick_model, +) -> None: + """`lite configure` with no agent named: ask which agents to wire and which model to pin.""" + targets: Final = pick_targets() + if _CLAUDE_TARGET not in targets: + return + credential, listed = _start(ctx, None) + _apply_claude(ctx, credential, listed, pick_model(listed)) + + +@click.group(name="configure", invoke_without_command=True) +@click.pass_context +def configure_group(ctx: click.Context) -> None: + """Persistently route a coding agent through your LiteLLM proxy. + + With no agent named, asks which agents to wire and which proxy model to pin. + """ + if ctx.invoked_subcommand is not None: + return + if not sys.stdin.isatty(): + raise click.ClickException( + "`lite configure` asks questions, so it needs a terminal. Non-interactively, run " + "`lite configure claude --api-key --model `." + ) + interactive_configure(ctx) + + +@click.group(name="unconfigure") +def unconfigure_group() -> None: + """Undo `lite configure` for a coding agent.""" + + +@configure_group.command(name="claude") +@click.option( + "--api-key", + "api_key", + default=None, + help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / " + "LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.", +) +@click.option("--model", default=None, help=_MODEL_OPTION_HELP) +@click.pass_context +def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: + """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. + + Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a + static token, or your `lite login` through apiKeyHelper), and gateway model discovery so + /model lists the proxy's models; --model picks the one Claude Code starts on. Every other + setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. + Assumes the proxy is already running. + """ + credential, listed = _start(ctx, api_key) + _apply_claude(ctx, credential, listed, model) + + +@unconfigure_group.command(name="claude") +def unconfigure_claude() -> None: + """Return Claude Code's settings to what they were before `lite configure claude`. + + Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are + put back; anything you changed since is left as it is and named in the output. + """ + try: + outcome: Final = unconfigure_claude_settings(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + _report_unconfigure(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, outcome) + + +def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None: + """Say what unconfigure did, naming only keys whose value it changed.""" + if outcome.file_removed: + click.echo( + f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys." + ) + elif outcome.restored: + click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.") + else: + click.echo(f"Nothing in {settings_path} was still ours to restore.") + if outcome.kept: + click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.") + if outcome.withheld: + click.echo( + "Left removed, since the file now points at a different server than they were issued for: " + + "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld) + + f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` " + "again to put them back, or delete that file to drop them." + ) + + +__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 7b0c1970c4e..70c89a853e3 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -10,6 +10,7 @@ import os import tempfile from collections.abc import Callable, Mapping from dataclasses import dataclass +from enum import StrEnum from pathlib import Path from types import MappingProxyType from typing import Final @@ -20,11 +21,28 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" PI_PROVIDER_NAME: Final = "litellm" LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" +_REJECTED_STATUSES: Final = frozenset((401, 403)) + + +class ListingFailure(StrEnum): + """Why a proxy could not be listed, decided once where the HTTP outcome is classified. + + `unreachable` means no response at all; the other kinds prove the proxy answered, so callers + must not suggest checking whether it is running. + """ + + UNREACHABLE = "unreachable" + REJECTED = "rejected" + BAD_BODY = "bad_body" + EMPTY = "empty" + OTHER = "other" @dataclass(frozen=True, slots=True) class PiSyncError: message: str + status: int | None = None + kind: ListingFailure | None = None @dataclass(frozen=True, slots=True) @@ -65,16 +83,20 @@ def fetch_model_ids( timeout=10, ) except requests.RequestException as e: - return PiSyncError(f"Could not list models from the proxy: {e}") + return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE) if resp.status_code != 200: - return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.") + return PiSyncError( + f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.", + resp.status_code, + ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER, + ) try: listing: Final = _ModelList.model_validate(resp.json()) except (ValueError, ValidationError) as e: - return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}") + return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY) ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) if not ids: - return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.") + return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) return ids @@ -200,6 +222,7 @@ __all__ = ( "LITELLM_PROXY_API_KEY_ENV", "PI_CONFIG_DIR_ENV", "PI_PROVIDER_NAME", + "ListingFailure", "ModelLimits", "PiSyncError", "fetch_model_ids", diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index b7c02866d6f..ffece87ab83 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_ from .claude_settings import ( BACKUP_PATH, CLAUDE_SETTINGS_PATH, + ApiKeyHelper, ClaudeSettingsError, load_json_or_empty, merge_claude_settings, @@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool: return token_data is not None and token_data.get("refresh_token") is not None -def _ensure_fresh_login(ctx: click.Context) -> None: +def ensure_fresh_login(ctx: click.Context) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"].rstrip("/") vault: Final = context_secret_vault(ctx) @@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None: click.echo("No fresh LiteLLM login found for this proxy; starting login...") ctx.invoke(login, pkce=pkce) if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault): - raise UpError("Login did not produce a usable token; cannot start `lite up`.") + raise UpError("Login did not produce a usable token.") def _restore_and_report() -> None: @@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None: base_url: Final = ctx.obj["base_url"] try: - _ensure_fresh_login(ctx) + ensure_fresh_login(ctx) api_key: Final = resolve_api_key(ctx) verify_proxy_key(base_url, api_key) @@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None: ) CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) - merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper) + merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper)) with open(CLAUDE_SETTINGS_PATH, "w") as f: json.dump(merged, f, indent=2) except (AgentRunError, ClaudeSettingsError) as e: diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index eae1b0f5bc9..b0e81a222c0 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names +from .commands.configure import configure_group, unconfigure_group from .commands.credentials import credentials from .commands.debug import debug from .commands.encryption import encryption @@ -162,6 +163,9 @@ cli.add_command(model_groups) # Add the autoroute command group (QA auto-routing against your real proxy) cli.add_command(autoroute_group, name="autoroute") cli.add_command(config_commands) +# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key) +cli.add_command(configure_group) +cli.add_command(unconfigure_group) if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 4a3b3ef22c4..77742ea9f9f 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -159,6 +159,10 @@ class TestUpCommand: assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert "apiKeyHelper" not in captured["settings"] + # The ephemeral proxy serves only the autorouter, so a starting model left by + # `lite configure claude --model` or a user pin would 400 on the first message. + assert captured["settings"]["model"] == "autorouter" + assert captured["settings"]["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" assert captured["settings_mode"] == 0o600 assert terminate_calls == [99999] diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py deleted file mode 100644 index 87a33c79a79..00000000000 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ /dev/null @@ -1,63 +0,0 @@ -from litellm.proxy.client.cli.commands.autoroute.settings import ( - ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, - merge_claude_settings_static_token, -) - - -def test_preserves_unrelated_top_level_keys(): - merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") - assert merged["theme"] == "dark" - - -def test_preserves_unrelated_env_keys(): - settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["SOME_OTHER_VAR"] == "value" - - -def test_sets_base_url_and_auth_token(): - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") - assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" - assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" - - -def test_preserves_existing_tool_search(): - settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" - - -def test_drops_stray_api_key(): - settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "ANTHROPIC_API_KEY" not in merged["env"] - - -def test_removes_existing_api_key_helper(): - settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "apiKeyHelper" not in merged - - -def test_does_not_mutate_input(): - settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - - -def test_forces_all_claude_code_default_model_tiers_to_the_autorouter(): - # A bare "*" model_name deployment looks like the obvious way to catch every request - # regardless of which model Claude Code thinks it's using, but Router's auto-router - # registry is keyed by the literal requested model string with no wildcard resolution - # (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude - # Code's own tiers hit the auto-router is to override the env vars it reads per tier. - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc") - for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: - assert merged["env"][key] == "autorouter" - - -def test_overrides_a_preexisting_default_model_env_var(): - settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 821323e722c..7d5cb61c062 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -27,6 +27,7 @@ from litellm.proxy.client.cli.commands.auth import ( print_token, whoami, ) +from litellm.proxy.client.cli.commands import auth as auth_module from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner @@ -1398,8 +1399,10 @@ class TestLoginConfigClaude: def setup_method(self): self.runner = CliRunner() - def _run_login(self, tmp_path, args, base_url="https://test.example.com"): + def _run_login(self, tmp_path, monkeypatch, args, base_url="https://test.example.com"): settings_path = tmp_path / "claude" / "settings.json" + monkeypatch.setattr(auth_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(auth_module, "CONFIGURE_STATE_PATH", tmp_path / "claude_configure_state.json") backup_path = tmp_path / "claude_settings_backup.json" poll_response = Mock() poll_response.status_code = 200 @@ -1416,7 +1419,6 @@ class TestLoginConfigClaude: patch("requests.get", return_value=poll_response), patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), patch("litellm.proxy.client.cli.interface.show_commands"), - patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), patch( "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),), @@ -1429,16 +1431,16 @@ class TestLoginConfigClaude: result = self.runner.invoke(login, args, obj={"base_url": base_url}) return result, settings_path, backup_path - def test_default_login_does_not_touch_claude_settings(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, []) + def test_default_login_does_not_touch_claude_settings(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, []) assert result.exit_code == 0 assert "Login successful!" in result.output assert not settings_path.exists() assert "Configured Claude Code" not in result.output - def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) @@ -1446,25 +1448,43 @@ class TestLoginConfigClaude: assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" assert "Configured Claude Code" in result.output + assert "pins a proxy model for every tier" not in result.output + assert "the model Claude Code starts on" in result.output - def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path): + def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}})) - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" assert written["env"]["KEEP"] == "me" - def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path): + def test_refuses_before_logging_in_while_lite_up_holds_the_settings(self, tmp_path, monkeypatch): + # The local precondition comes first: no browser, no token stored, no "Login successful!". + backup_path = tmp_path / "claude_settings_backup.json" + backup_path.write_text("{}") + monkeypatch.setattr(auth_module, "CLAUDE_SETTINGS_PATH", tmp_path / "claude" / "settings.json") + monkeypatch.setattr( + auth_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),) + ) + with patch("requests.post") as post, patch("webbrowser.open") as browser: + result = self.runner.invoke(login, ["--config-claude"], obj={"base_url": "https://test.example.com"}) + assert result.exit_code != 0 + assert "not logging in" in result.output and "lite down" in result.output + assert "Login successful!" not in result.output + post.assert_not_called() + browser.assert_not_called() + + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text("not json at all {{{") - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code != 0 assert "Login successful!" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index e5f2a9d95bd..9072fe26b00 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -1,4 +1,5 @@ import json +import os import shlex import stat import time @@ -9,14 +10,25 @@ from click.testing import CliRunner from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord from litellm.proxy.client.cli import cli +from litellm.litellm_core_utils.private_json import commit_staged_json from litellm.proxy.client.cli.commands.claude_settings import ( + ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, BACKUP_PATH, + OWNED_ENV_KEYS, + OWNED_TOP_LEVEL_KEYS, SETTINGS_FILE_OWNERS, + ApiKeyHelper, ClaudeSettingsError, + KeepModel, SettingsFileOwner, + StartOn, + StaticToken, + UnpinModel, + configure_claude_settings, + merge_claude_settings, resolve_api_key_helper, - write_claude_settings, + unconfigure_claude_settings, ) @@ -24,6 +36,7 @@ def _owners(*backup_paths): """Stand-in owners for the real `lite up` / `lite autoroute up` registry.""" return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths) + CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" @@ -97,17 +110,28 @@ def lite_on_path(): yield -class TestWriteClaudeSettings: +def _helper_configure(base_url, settings_path, owners, state_path=None): + """`lite login --config-claude`'s shape: the login credential behind apiKeyHelper, no pinned model.""" + state = state_path if state_path is not None else settings_path.parent.parent / "state.json" + root = base_url.rstrip("/") + configure_claude_settings( + root, ApiKeyHelper(resolve_api_key_helper(root)), KeepModel(), settings_path, state, owners + ) + + +class TestConfigureWithTheLoginHelper: def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path): settings_path, backup_path = paths assert not settings_path.parent.exists() - write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com/", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + assert "model" not in written def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): settings_path, backup_path = paths @@ -123,7 +147,7 @@ class TestWriteClaudeSettings: ) ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" @@ -135,26 +159,31 @@ class TestWriteClaudeSettings: def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path): settings_path, backup_path = paths - write_claude_settings("https://first.example.com", settings_path, _owners(backup_path)) - write_claude_settings("https://second.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://first.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://second.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com" assert "second.example.com" in written["apiKeyHelper"] assert "first.example.com" not in written["apiKeyHelper"] - def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path): + def test_drops_stray_static_credentials_so_the_helper_token_wins(self, paths, lite_on_path): + # Claude Code prefers ANTHROPIC_AUTH_TOKEN over apiKeyHelper, so a virtual key left behind + # by an earlier `lite configure claude --api-key` would silently keep winning. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) - settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}})) + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked", "ANTHROPIC_AUTH_TOKEN": "sk-old"}}) + ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"] + env = json.loads(settings_path.read_text())["env"] + assert "ANTHROPIC_API_KEY" not in env and "ANTHROPIC_AUTH_TOKEN" not in env def test_written_file_is_owner_only(self, paths, lite_on_path): settings_path, backup_path = paths - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600 def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path): @@ -162,7 +191,7 @@ class TestWriteClaudeSettings: backup_path.write_text("{}") with pytest.raises(ClaudeSettingsError, match="lite down"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() @@ -172,7 +201,7 @@ class TestWriteClaudeSettings: settings_path.write_text("not json at all {{{") with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert settings_path.read_text() == "not json at all {{{" @@ -180,7 +209,7 @@ class TestWriteClaudeSettings: settings_path, backup_path = paths with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() @@ -196,7 +225,7 @@ class TestWriteClaudeSettings: settings_path.write_bytes(b'{"theme": "\xff\xfe"}') with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path): """An unreadable settings file must not surface as "Authentication failed". @@ -210,16 +239,18 @@ class TestWriteClaudeSettings: settings_path.mkdir() with pytest.raises(ClaudeSettingsError, match="Could not read"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path): settings_path, backup_path = paths - with patch( - f"{CLAUDE_SETTINGS_MODULE}.write_private_json", - side_effect=OSError("Read-only file system"), - ): - with pytest.raises(ClaudeSettingsError, match="Read-only file system"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + settings_path.parent.mkdir(parents=True) + settings_path.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + finally: + settings_path.parent.chmod(0o700) + assert not settings_path.exists() class TestApiKeyHelperIsActuallyInvocable: @@ -303,7 +334,7 @@ class TestConflictingOwnersOfTheSettingsFile: backup.write_text("{}") stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command) with pytest.raises(ClaudeSettingsError, match="currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (stand_in,)) + _helper_configure("https://proxy.example.com", settings_path, (stand_in,)) backup.unlink() assert not settings_path.exists() @@ -314,9 +345,9 @@ class TestConflictingOwnersOfTheSettingsFile: autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down") with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) def test_the_registry_matches_the_paths_the_commands_actually_use(self): """A second definition of the autoroute dir must not drift from this one.""" @@ -341,7 +372,7 @@ class TestDoesNotDestroyUserOwnedStructure: link.parent.mkdir() link.symlink_to(real) - write_claude_settings("https://proxy.example.com", link, ()) + _helper_configure("https://proxy.example.com", link, ()) assert link.is_symlink() assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" @@ -354,6 +385,456 @@ class TestDoesNotDestroyUserOwnedStructure: settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"})) with pytest.raises(ClaudeSettingsError, match="non-object"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert json.loads(settings_path.read_text())["env"] == "not-an-object" + + +class TestMergeClaudeSettings: + """One merge for every way Claude Code gets wired: `lite up`, `lite login --config-claude`, + `lite configure claude` and `lite autoroute up`.""" + + def test_a_static_token_lands_in_env_and_the_helper_slot_is_cleared(self): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "leaked"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000/", StaticToken("token-abc")) + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert "ANTHROPIC_API_KEY" not in merged["env"] + assert "apiKeyHelper" not in merged + assert "model" not in merged + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_a_helper_lands_top_level_and_the_static_slots_are_cleared(self): + settings = {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-old", "ANTHROPIC_API_KEY": "leaked"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", ApiKeyHelper("lite auth print-token")) + assert merged["apiKeyHelper"] == "lite auth print-token" + assert "ANTHROPIC_AUTH_TOKEN" not in merged["env"] and "ANTHROPIC_API_KEY" not in merged["env"] + + def test_keeps_existing_switch_values_and_unrelated_keys_without_mutating_the_input(self): + settings = {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) + assert merged["theme"] == "dark" + assert merged["env"]["SOME_OTHER_VAR"] == "value" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" + assert settings == {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + + def test_a_default_model_sets_only_the_row_claude_code_starts_on(self): + merged = merge_claude_settings( + {}, "http://127.0.0.1:4000", StaticToken("token-abc"), default_model="claude-auto" + ) + assert merged["model"] == "claude-auto" + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_a_tier_model_forces_every_claude_code_tier_as_autoroute_needs(self): + # Router's auto-router registry is keyed by the literal requested model string with no + # wildcard resolution, so `lite autoroute up` overrides the env var each tier reads. + settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} + merged = merge_claude_settings( + settings, "http://127.0.0.1:4000", StaticToken("token-abc"), tier_model="autorouter" + ) + assert {merged["env"][key] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS} == {"autorouter"} + assert "model" not in merged + + def test_touches_exactly_the_declared_owned_keys(self): + # The receipt and unconfigure restore exactly OWNED_*_KEYS, so a key the merge writes outside + # that table would be written by configure and never undone. + settings = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "old", "ENABLE_TOOL_SEARCH": "false"}, + "apiKeyHelper": "old-helper", + "model": "old-model", + } + for credential in (StaticToken("token-abc"), ApiKeyHelper("helper")): + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", credential, default_model="claude-auto") + changed_top_level = {key for key in set(settings) | set(merged) if settings.get(key) != merged.get(key)} + assert changed_top_level - {"env"} <= set(OWNED_TOP_LEVEL_KEYS) + changed_env = { + key + for key in set(settings["env"]) | set(merged["env"]) + if settings["env"].get(key) != merged["env"].get(key) + } + assert changed_env <= set(OWNED_ENV_KEYS) + assert merged["permissions"] == {"allow": ["Bash"]} + assert merged["env"]["KEEP_ME"] == "1" + + +PROXY = "http://127.0.0.1:4000" +ANTHROPIC = "https://api.anthropic.com" +HELPER = ApiKeyHelper("lite auth print-token") +ORIGINAL = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC}, + "apiKeyHelper": "/usr/local/bin/lite auth print-token", + "model": "claude-opus-5", +} + + +def _set(path, value): + """A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file.""" + + def edit(settings): + section, _, key = path.rpartition(".") + container = settings.setdefault(section, {}) if section else settings + if value is _ABSENT: + container.pop(key, None) + else: + container[key] = value + return settings + + return edit + + +_ABSENT = object() + + +class _Rig: + """One settings file plus receipt under tmp_path, driven through the public functions only.""" + + def __init__(self, tmp_path, initial): + self.settings = tmp_path / "claude" / "settings.json" + self.state = tmp_path / "state" / "claude_configure_state.json" + if initial is not None: + self.settings.parent.mkdir(parents=True) + self.settings.write_text(json.dumps(initial)) + + def read(self): + return json.loads(self.settings.read_text()) if self.settings.exists() else None + + def configure(self, credential=StaticToken("sk-virtual-key"), model=StartOn("claude-auto"), **kwargs): + configure_claude_settings(PROXY, credential, model, self.settings, self.state, (), **kwargs) + + def edit(self, *edits): + settings = self.read() + for apply in edits: + settings = apply(settings) + self.settings.write_text(json.dumps(settings)) + + def unconfigure(self): + return unconfigure_claude_settings(self.settings, self.state, ()) + + +# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure +# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used +# to be one test each; the receipt's rules are what make them all come out right. +UNDO_SCENARIOS = { + "plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}), + "no file before": (None, [], None, {"file_removed": True}), + "no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}), + "null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}), + "empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}), + "user edits stay and are named": ( + ORIGINAL, + [_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()}, + ), + "user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}), + "user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}), + "user removed our key: neither restored nor kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)], + ORIGINAL, + {"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()}, + ), + "restored names only what changed": ( + {"model": "claude-opus-5"}, + [], + {"model": "claude-opus-5"}, + { + "restored": { + "env.ANTHROPIC_BASE_URL", + "env.ENABLE_TOOL_SEARCH", + "env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", + "apiKeyHelper", + }, + "kept": (), + }, + {"credential": HELPER, "model": KeepModel()}, + ), + "repeat across credential kinds keeps the first snapshot": ( + ORIGINAL, + [ + {"credential": HELPER, "model": UnpinModel()}, + {"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")}, + ], + ORIGINAL, + {}, + ), + "repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}), + "repeat without a model lets go of our pin, user had one": ( + {"model": "claude-opus-5"}, + [{"model": UnpinModel()}], + {"model": "claude-opus-5"}, + {}, + ), + "re-login keeps our pin": (None, [{"credential": HELPER, "model": KeepModel()}], None, {"file_removed": True}), + "edit between configures survives an unpin repeat": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": HELPER, "model": UnpinModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures survives a re-login": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": HELPER, "model": KeepModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures: a same-model repeat displaces it, so it is what comes back": ( + ORIGINAL, + [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": HELPER}], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}}, + ), + "base URL changed since: credentials withheld, receipt kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT}, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL"}, + "receipt_kept": True, + }, + ), + "base URL changed and back: judged against the URL the restored file holds": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)], + ORIGINAL, + {"withheld": ()}, + ), + "credential captured beside no URL goes back only beside no URL": ( + {"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}}, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}}, + { + "withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")}, + "receipt_kept": True, + }, + ), + "restored document empty while a credential is withheld: file goes, receipt stays": ( + None, + [ + _set("env.ANTHROPIC_API_KEY", "sk-user"), + {"credential": HELPER, "model": KeepModel()}, + _set("env.ANTHROPIC_BASE_URL", _ABSENT), + ], + None, + {"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True}, + {"credential": HELPER, "model": KeepModel()}, + ), + "a credential the user changed is kept, never also withheld": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")], + { + **ORIGINAL, + "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, + "apiKeyHelper": "/opt/mine/helper", + }, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"}, + "receipt_kept": True, + }, + ), +} + + +def _expected_file(expected): + if expected is None: + return None + return {k: v for k, v in expected.items() if v is not _ABSENT} + + +class TestConfigureAndUnconfigure: + """`configure_claude_settings` records how to undo itself; `unconfigure_claude_settings` undoes only that.""" + + @pytest.mark.parametrize("scenario", UNDO_SCENARIOS.values(), ids=UNDO_SCENARIOS.keys()) + def test_undo_matrix(self, tmp_path, scenario): + initial, steps, expected, outcome_expectations, *first = scenario + rig = _Rig(tmp_path, initial) + rig.configure(**(first[0] if first else {})) + for step in steps: + if isinstance(step, dict): + rig.configure(**step) + elif rig.settings.exists() and step(json.loads(rig.settings.read_text())) is None: + rig.settings.unlink() + else: + rig.edit(step) + + outcome = rig.unconfigure() + + assert rig.read() == _expected_file(expected) + assert rig.state.exists() == outcome_expectations.get("receipt_kept", False) + for field, want in outcome_expectations.items(): + if field == "withheld": + assert {(item.key, item.endpoint) for item in outcome.withheld} == set(want) + elif field == "not_restored": + assert not set(want) & set(outcome.restored) and not set(want) & set(outcome.kept) + elif field == "restored_includes": + assert set(want) <= set(outcome.restored) + elif field in ("restored", "kept"): + assert set(getattr(outcome, field)) == set(want) + elif field != "receipt_kept": + assert getattr(outcome, field) == want + assert not {item.key for item in outcome.withheld} & set(outcome.kept) + + def test_configure_writes_owner_only_and_the_receipt_never_holds_the_key(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure(credential=StaticToken("sk-virtual-key-never-on-disk-twice")) + configured = rig.read() + assert configured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key-never-on-disk-twice" + assert configured["env"]["ANTHROPIC_BASE_URL"] == PROXY and configured["model"] == "claude-auto" + assert "ANTHROPIC_API_KEY" not in configured["env"] and "apiKeyHelper" not in configured + assert stat.S_IMODE(rig.settings.stat().st_mode) == 0o600 == stat.S_IMODE(rig.state.stat().st_mode) + assert "sk-virtual-key-never-on-disk-twice" not in rig.state.read_text() + + def test_withheld_credentials_come_back_once_the_url_points_at_their_server_again(self, tmp_path): + # The kept receipt owns only the withheld slots: the second unconfigure restores exactly those. + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")) + rig.unconfigure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", ANTHROPIC), _set("theme", "light")) + outcome = rig.unconfigure() + assert rig.read() == {**ORIGINAL, "theme": "light"} + assert set(outcome.restored) == {"env.ANTHROPIC_API_KEY", "apiKeyHelper"} + assert outcome.kept == () and outcome.withheld == () and not rig.state.exists() + + @pytest.mark.parametrize( + ("path", "value", "repeat_credential"), + [ + ("env.ANTHROPIC_API_KEY", "sk-user-added-later", HELPER), + ("env.ANTHROPIC_AUTH_TOKEN", "sk-users-own-token", HELPER), + ("apiKeyHelper", "/opt/mine/helper", StaticToken("sk-rotated")), + ], + ids=["user-adds-api-key", "user-replaces-our-token", "user-sets-own-helper"], + ) + def test_a_credential_the_user_set_between_two_configures_is_what_comes_back( + self, tmp_path, path, value, repeat_credential + ): + # The repeat's merge clears the slot, so the displaced value is snapshotted and is what returns; + # it was set while the file pointed at the proxy, so it returns once the file points there again. + rig = _Rig(tmp_path, {"theme": "dark"}) + rig.configure(credential=HELPER, model=KeepModel()) + rig.edit(_set(path, value)) + rig.configure(credential=repeat_credential, model=KeepModel()) + assert not _lookup(rig.read(), path) + + outcome = rig.unconfigure() + assert [(item.key, item.endpoint) for item in outcome.withheld] == [(path, PROXY)] + assert rig.read() == {"theme": "dark"} and rig.state.exists() + rig.settings.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_BASE_URL": PROXY}})) + outcome = rig.unconfigure() + assert _lookup(rig.read(), path) == value + assert outcome.restored == (path,) and outcome.withheld == () and not rig.state.exists() + + def test_a_receipt_commit_that_fails_leaves_no_staged_token_behind(self, tmp_path): + rig = _Rig(tmp_path, {}) + + def commit_receipt_fails(staged, path): + if path == str(rig.state): + os.unlink(staged) + raise OSError("receipt rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match=r"Could not write .*receipt rename failed"): + rig.configure(credential=StaticToken("sk-never-left-in-a-temp-file"), commit=commit_receipt_fails) + assert not list(rig.settings.parent.glob(".tmp-*")) and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read() == {} and not rig.state.exists() + + @pytest.mark.parametrize("configured_before", [False, True], ids=["first-configure", "repeat-configure"]) + def test_a_settings_commit_that_fails_after_the_receipt_landed_puts_the_receipt_back( + self, tmp_path, configured_before + ): + # The two renames are not atomic: a settings rename that fails after the receipt landed must + # not leave a receipt describing settings that were never written. + rig = _Rig(tmp_path, ORIGINAL) + if configured_before: + rig.configure() + receipt_before = rig.state.read_text() if configured_before else None + settings_before = rig.settings.read_text() + + def commit_settings_fails(staged, path): + if path == str(rig.settings): + os.unlink(staged) + raise OSError("rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match="rename failed"): + rig.configure(credential=StaticToken("sk-rotated"), commit=commit_settings_fails) + assert rig.settings.read_text() == settings_before + assert (rig.state.read_text() if rig.state.exists() else None) == receipt_before + if configured_before: + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_a_failed_repeat_configure_leaves_the_earlier_undo_intact(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + receipt_before = rig.state.read_text() + rig.settings.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + rig.configure(credential=StaticToken("sk-rotated")) + finally: + rig.settings.parent.chmod(0o700) + assert rig.state.read_text() == receipt_before and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read()["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_unconfigure_reports_a_receipt_it_cannot_remove_as_a_settings_error(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.state.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not remove"): + rig.unconfigure() + finally: + rig.state.parent.chmod(0o700) + + def test_configure_writes_through_a_symlinked_settings_file(self, tmp_path): + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text(json.dumps({"theme": "dark"})) + link = tmp_path / "settings.json" + link.symlink_to(target) + configure_claude_settings(PROXY, StaticToken("sk-virtual-key"), UnpinModel(), link, tmp_path / "state.json", ()) + assert link.is_symlink() + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + + @pytest.mark.parametrize("operation", ["configure", "unconfigure"]) + def test_refuses_while_a_temporary_owner_holds_a_backup(self, paths, tmp_path, operation): + settings_path, backup_path = paths + backup_path.write_text("{}") + owners = _owners(backup_path) + state = tmp_path / "state.json" + attempt = ( + (lambda: configure_claude_settings(PROXY, StaticToken("k"), UnpinModel(), settings_path, state, owners)) + if operation == "configure" + else (lambda: unconfigure_claude_settings(settings_path, state, owners)) + ) + with pytest.raises(ClaudeSettingsError, match="lite down"): + attempt() + assert not settings_path.exists() + + def test_unconfigure_without_a_receipt_is_an_error_not_a_silent_no_op(self, tmp_path): + with pytest.raises(ClaudeSettingsError, match="nothing to undo"): + _Rig(tmp_path, None).unconfigure() + + +def _lookup(settings, path): + section, _, key = path.rpartition(".") + return (settings.get(section) or {}).get(key) if section else settings.get(key) diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py new file mode 100644 index 00000000000..1631ba7f49e --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -0,0 +1,331 @@ +import json +import os +import stat + +import click +import pytest +import requests +import responses +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands import configure as configure_module +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner +from litellm.proxy.client.cli.commands.configure import configure_claude, configure_group, interactive_configure + +PROXY = "http://proxy.test:4000" +VALID_KEY = "sk-virtual-key" +LISTED_MODELS = ("claude-auto", "gpt-5.6-luna") + + +def _mock_models(): + responses.get( + f"{PROXY}/v1/models", + json={"data": [{"id": model, "object": "model"} for model in LISTED_MODELS]}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}"})], + ) + responses.get(f"{PROXY}/v1/models", status=401) + + +@pytest.fixture +def paths(monkeypatch, tmp_path): + settings_path = tmp_path / "claude" / "settings.json" + state_path = tmp_path / "litellm" / "claude_configure_state.json" + monkeypatch.setattr(configure_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(configure_module, "CONFIGURE_STATE_PATH", state_path) + return settings_path, state_path + + +@pytest.fixture +def lite_on_path(monkeypatch, tmp_path): + """A real `lite` executable on PATH, so the apiKeyHelper command resolves without patching.""" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + lite = bin_dir / "lite" + lite.write_text("#!/bin/sh\nexit 0\n") + lite.chmod(lite.stat().st_mode | stat.S_IXUSR) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ.get('PATH', '')}") + return str(lite) + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def lite_up_backup(monkeypatch, tmp_path): + """A `lite up` session holding its backup, the local precondition every settings write refuses on.""" + backup = tmp_path / "claude_settings_backup.json" + backup.write_text("{}") + monkeypatch.setattr(configure_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup, "lite up", "lite down"),)) + return backup + + +def _configure(runner, *args): + return runner.invoke(cli, ["--base-url", PROXY, "configure", "claude", *args]) + + +class TestConfigureClaudeWithAVirtualKey: + @responses.activate + def test_writes_settings_and_reports_without_echoing_the_key(self, runner, paths): + _mock_models() + settings_path, state_path = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert written["model"] == "claude-auto" + assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in written["env"] + assert state_path.exists() + assert VALID_KEY not in result.output + assert "Starting model: claude-auto" in result.output + assert "1 of the proxy's 2 models" in result.output + assert "lite unconfigure claude" in result.output + assert len(responses.calls) == 1 + + @responses.activate + def test_takes_the_key_from_the_global_option_and_keeps_claude_codes_default(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = runner.invoke(cli, ["--base-url", PROXY, "--api-key", VALID_KEY, "configure", "claude"]) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert "model" not in written + assert "Starting model: not pinned" in result.output + + @responses.activate + def test_refuses_a_model_the_proxy_does_not_list(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-nope") + assert result.exit_code != 0 + assert "'claude-nope' is not served" in result.output + assert "claude-auto, gpt-5.6-luna" in result.output + assert not settings_path.exists() + + @responses.activate + def test_refuses_a_key_the_proxy_rejects(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", "sk-wrong") + assert result.exit_code != 0 + assert "rejected your key (HTTP 401)" in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize( + ("mock", "expected", "unexpected"), + [ + ( + lambda: responses.get(f"{PROXY}/v1/models", body=requests.ConnectionError("refused")), + "Is the proxy at", + "answered", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", status=500), + "The proxy at http://proxy.test:4000 answered", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", body="not json"), + "answered, so check that it is a LiteLLM proxy", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", json={"data": []}), + "Claude Code would have nothing to run", + "Is the proxy at", + ), + ], + ids=["unreachable", "http-500", "non-json-body", "empty-list"], + ) + def test_the_listing_hint_matches_how_the_listing_failed(self, runner, paths, mock, expected, unexpected): + # Only a proxy that never answered gets the "is it running" question; a 500, a non-JSON body or an + # empty list prove it is up, and the hint says so instead. + mock() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code != 0 + assert expected in result.output and unexpected not in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize("entry", ["virtual-key", "login", "interactive"]) + def test_refuses_while_lite_up_holds_a_backup_before_any_login_or_request( + self, runner, paths, monkeypatch, lite_up_backup, entry + ): + _mock_models() + + def login_must_not_run(ctx): + raise AssertionError("the local precondition must be checked before a login is attempted") + + monkeypatch.setattr(configure_module, "ensure_fresh_login", login_must_not_run) + if entry == "interactive": + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": None}) + with pytest.raises(click.ClickException, match="lite down"): + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None) + else: + args = ["--api-key", VALID_KEY] if entry == "virtual-key" else [] + result = runner.invoke(configure_claude, args, obj={"base_url": PROXY, "api_key": None}) + assert result.exit_code != 0 and "lite down" in result.output + assert len(responses.calls) == 0 + assert not paths[0].exists() + + @responses.activate + def test_says_so_when_the_key_is_written_through_a_symlink(self, runner, paths, tmp_path): + _mock_models() + settings_path, _ = paths + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text("{}") + settings_path.parent.mkdir(parents=True) + settings_path.symlink_to(target) + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "keep it out of version control" in result.output + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + + +class TestConfigureClaudeWithTheLogin: + def _stored_login(self, monkeypatch): + monkeypatch.setattr(configure_module, "ensure_fresh_login", lambda ctx: None) + monkeypatch.setattr(configure_module, "get_stored_api_key", lambda expected_base_url, vault: VALID_KEY) + + @responses.activate + def test_uses_the_login_through_the_helper_and_writes_no_secret(self, runner, paths, monkeypatch, lite_on_path): + _mock_models() + self._stored_login(monkeypatch) + settings_path, _ = paths + result = runner.invoke( + configure_claude, + ["--model", "claude-auto"], + obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": True}, + ) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["apiKeyHelper"] == f"{lite_on_path} --base-url {PROXY} auth print-token" + assert "ANTHROPIC_AUTH_TOKEN" not in written["env"] + assert written["model"] == "claude-auto" + assert VALID_KEY not in settings_path.read_text() + assert "read through apiKeyHelper" in result.output + + @responses.activate + def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths, monkeypatch, lite_on_path): + _mock_models() + self._stored_login(monkeypatch) + settings_path, _ = paths + result = runner.invoke( + configure_claude, + ["--api-key", VALID_KEY], + obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, + ) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY and "apiKeyHelper" not in written + + +class TestInteractiveConfigure: + @responses.activate + def test_asks_for_targets_and_a_starting_model_then_configures(self, paths): + _mock_models() + settings_path, _ = paths + asked = {} + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "claude-auto" + + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == LISTED_MODELS + assert json.loads(settings_path.read_text())["model"] == "claude-auto" + + def test_does_nothing_when_claude_code_is_not_picked(self, paths): + settings_path, _ = paths + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda listed: None) + assert not settings_path.exists() + + def test_bare_configure_without_a_terminal_names_the_non_interactive_command(self, runner, paths): + result = runner.invoke(cli, ["--base-url", PROXY, "configure"]) + assert result.exit_code != 0 + assert "lite configure claude --api-key" in result.output + + +class TestUnconfigureClaude: + @responses.activate + def test_restores_the_original_file_and_removes_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + original = {"theme": "dark", "model": "claude-opus-5"} + settings_path.write_text(json.dumps(original)) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == original + assert not state_path.exists() + assert "Restored in" in result.output and "model" in result.output + assert "ANTHROPIC_API_KEY" not in result.output, "a key that never existed was not restored" + + @responses.activate + def test_a_file_only_configure_created_is_reported_removed_not_restored(self, runner, paths): + _mock_models() + settings_path, _ = paths + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert not settings_path.exists() + assert "No settings file remains" in result.output and "Restored" not in result.output + + @responses.activate + def test_says_when_nothing_was_still_ours_and_names_what_it_kept(self, runner, paths): + _mock_models() + settings_path, _ = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark"})) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()} + edited["model"] = "mine" + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "Nothing in" in result.output and "was still ours to restore" in result.output + assert "Left as you changed them since:" in result.output and "model" in result.output + + @responses.activate + def test_names_the_server_a_withheld_credential_was_captured_with_and_keeps_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://api.anthropic.com", "ANTHROPIC_API_KEY": "sk-ant"}}) + ) + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"]["ANTHROPIC_BASE_URL"] = "http://other-proxy:4000" + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "env.ANTHROPIC_API_KEY (captured with https://api.anthropic.com)" in result.output + assert str(state_path) in result.output and state_path.exists() + assert "sk-ant" not in result.output + + def test_refuses_while_lite_up_holds_a_backup(self, runner, paths, lite_up_backup): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 and "lite down" in result.output + + def test_without_a_receipt_it_fails_loudly(self, runner, paths): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 + assert "nothing to undo" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 68c0ac70064..1c2da514ee2 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -4,9 +4,11 @@ import stat from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import pytest import requests from litellm.proxy.client.cli.commands.pi import ( + ListingFailure, ModelLimits, PiSyncError, fetch_model_ids, @@ -28,6 +30,10 @@ class _FakeResponse: return self._payload +def _refused(*args, **kwargs): + raise requests.ConnectionError("refused") + + class TestFetchModelIds: def test_returns_ids_in_proxy_order_deduped(self): captured = {} @@ -53,9 +59,7 @@ class TestFetchModelIds: assert "Could not list models" in result.message def test_non_200_is_a_value(self): - result = fetch_model_ids( - "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500) - ) + result = fetch_model_ids("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)) assert isinstance(result, PiSyncError) assert "HTTP 500" in result.message @@ -75,6 +79,22 @@ class TestFetchModelIds: ) assert isinstance(result, PiSyncError) assert "no models" in result.message + assert result.kind is ListingFailure.EMPTY + + @pytest.mark.parametrize( + ("get", "kind"), + [ + (_refused, ListingFailure.UNREACHABLE), + (lambda *a, **k: _FakeResponse(401), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(403), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(500), ListingFailure.OTHER), + (lambda *a, **k: _FakeResponse(200), ListingFailure.BAD_BODY), + ], + ids=["unreachable", "401", "403", "500", "bad-body"], + ) + def test_the_failure_kind_is_decided_where_the_response_is_classified(self, get, kind): + result = fetch_model_ids("http://localhost:4000", "sk-key", get=get) + assert isinstance(result, PiSyncError) and result.kind is kind class TestFetchModelLimits: diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index aead1764b0e..b607b1a3db7 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -11,11 +11,11 @@ from click.testing import CliRunner from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError -from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.claude_settings import ApiKeyHelper, ClaudeSettingsError from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, - _ensure_fresh_login, + ensure_fresh_login, down, load_json_or_empty, merge_claude_settings, @@ -40,12 +40,12 @@ def _patch_paths(monkeypatch, tmp_path): class TestMergeClaudeSettings: def test_preserves_unrelated_top_level_keys(self): - merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper") + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["theme"] == "dark" def test_preserves_unrelated_env_keys(self): settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["SOME_OTHER_VAR"] == "value" def test_overrides_base_url_and_helper(self): @@ -53,7 +53,7 @@ class TestMergeClaudeSettings: "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, "apiKeyHelper": "old-helper", } - merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") + merged = merge_claude_settings(settings, "http://localhost:4000/", ApiKeyHelper("new-helper")) assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" @@ -61,21 +61,21 @@ class TestMergeClaudeSettings: def test_preserves_existing_gateway_model_discovery(self): settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert "ANTHROPIC_API_KEY" not in merged["env"] def test_works_from_empty_settings(self): - merged = merge_claude_settings({}, "http://localhost:4000", "helper") + merged = merge_claude_settings({}, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"] == { "ANTHROPIC_BASE_URL": "http://localhost:4000", "ENABLE_TOOL_SEARCH": "true", @@ -85,7 +85,7 @@ class TestMergeClaudeSettings: def test_does_not_mutate_input(self): settings = {"env": {"FOO": "bar"}} - merge_claude_settings(settings, "http://localhost:4000", "helper") + merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert settings == {"env": {"FOO": "bar"}} @@ -327,7 +327,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] @@ -339,7 +339,7 @@ class TestEnsureFreshLogin: monkeypatch, on_login=lambda: store.log_in({"key": "sk-b", "base_url": "http://proxy-b:4000"}, "sk-b") ) - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) assert login_calls == [("http://proxy-b:4000", False)] assert store.key_requests == ["http://proxy-b:4000", "http://proxy-b:4000"] @@ -353,7 +353,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in({"key": "sk-a", "base_url": "http://proxy-a:4000"}, "sk-a"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", False)] @@ -363,7 +363,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) with pytest.raises(UpError, match="Run `lite login` first"): - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) def test_trusts_a_pkce_credential_that_was_renewed_on_the_way_in(self, monkeypatch): """A --pkce key inside its freshness buffer is renewed by `get_stored_api_key`, so `lite up` @@ -377,7 +377,7 @@ class TestEnsureFreshLogin: ) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] assert store.key_requests == ["http://proxy-a:4000"] @@ -390,7 +390,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", True)] @@ -399,7 +399,7 @@ class TestEnsureFreshLogin: _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=-10), {}) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) def test_trusts_the_key_the_cli_group_already_resolved_instead_of_reading_the_token_file_again( self, monkeypatch @@ -409,7 +409,7 @@ class TestEnsureFreshLogin: store = _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=86_400), {}) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) assert login_calls == [] assert store.key_requests == [] @@ -423,7 +423,7 @@ class TestEnsureFreshLogin: ) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert store.key_requests == [] @@ -435,7 +435,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert login_calls == [("http://proxy-a:4000", True)] assert store.key_requests == ["http://proxy-a:4000"] From e34c4c8edc5fc6cd70fccf6a44c7765d920a6e28 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 18:09:28 -0700 Subject: [PATCH 155/157] fix(ui): scope shadow eval models to configured chat groups (#40488) Co-authored-by: Claude Code --- .../_components/ShadowEvalSection.test.tsx | 151 ++++++++++++++---- .../_components/ShadowEvalStartForm.tsx | 81 ++++------ .../hooks/models/useModels.test.ts | 103 +++++++++--- .../app/(dashboard)/hooks/models/useModels.ts | 61 ++++--- .../components/chat_ui/EndpointUtils.tsx | 35 +--- .../chat_ui/mode_endpoint_mapping.tsx | 17 ++ .../src/lib/autorouter_presets.test.ts | 13 ++ .../src/lib/autorouter_presets.ts | 13 +- 8 files changed, 317 insertions(+), 157 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 64363da9933..1f73671caae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -1,3 +1,4 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; @@ -14,7 +15,9 @@ vi.mock("./useShadowEval", () => ({ })); const authorizedRoleMock = vi.fn(() => ({ accessToken: "token", isViewOnly: false })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => authorizedRoleMock() })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ userId: "test-user-id", userRole: "Admin", ...authorizedRoleMock() }), +})); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useInfiniteKeys: vi.fn(() => ({ @@ -68,27 +71,33 @@ vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ })), })); -vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ +vi.mock("@/app/(dashboard)/hooks/models/useModels", async (importOriginal) => ({ + ...(await importOriginal()), useAutoRouters: vi.fn(() => ({ data: [ { model_name: "claude-auto", litellm_params: { model: "auto_router/claude-auto" } }, { model_name: "gpt-auto", litellm_params: { model: "auto_router/gpt-auto" } }, ], })), - usePlainModelGroups: vi.fn(() => new Set(["prod-claude"])), + usePlainModelGroups: vi.fn(() => new Set(["prod-claude", "prod-judge"])), + usePlainChatModelGroups: vi.fn(() => new Set(["prod-claude", "prod-judge"])), + usePlainChatModelDeployments: vi.fn(() => [ + { + model_name: "prod-judge", + litellm_params: { model: "anthropic/claude-sonnet-5" }, + model_info: { mode: "chat" }, + }, + ]), })); -vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ - useModelCostMap: vi.fn(() => ({ - data: { - "claude-sonnet-5": { litellm_provider: "anthropic", mode: "chat" }, - "gpt-4o": { litellm_provider: "openai", mode: "chat" }, - "gemini/gemini-2.5-pro": { litellm_provider: "gemini", mode: "chat" }, - "text-embedding-3-large": { litellm_provider: "openai", mode: "embedding" }, - }, - })), +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + modelInfoCall: vi.fn(), })); +import { usePlainChatModelGroups, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { modelInfoCall } from "@/components/networking"; + import ShadowEvalSection, { shadowedTargetLabel } from "./ShadowEvalSection"; import { useShadowEvalJob, @@ -107,7 +116,7 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ models: [], direction: "forward", baseline_model: null, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", shadow_percentage: 10, targets: [ { @@ -249,6 +258,85 @@ describe("ShadowEvalSection", () => { if (defaultKeysImpl) vi.mocked(useInfiniteKeys).mockImplementation(defaultKeysImpl); }); + it("labels only configured judge recommendations", async () => { + const user = userEvent.setup(); + mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Select a judge model")); + expect(screen.getByRole("option", { name: /prod-judge.*Recommended/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); + + await user.keyboard("{Escape}"); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + expect(screen.getByRole("option", { name: "prod-judge", exact: true })).toBeInTheDocument(); + expect(screen.queryByText("Recommended")).not.toBeInTheDocument(); + }); + + it("keeps custom models selectable through the real model hooks without widening chat choices to traffic filters", async () => { + const hooks = await vi.importActual( + "@/app/(dashboard)/hooks/models/useModels", + ); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const deployments = [ + { model_name: "custom-chat", litellm_params: { model: "openai/private-chat" } }, + { model_name: "custom-judge", litellm_params: { model: "openai/private-judge" }, model_info: { mode: null } }, + { + model_name: "embedding", + litellm_params: { model: "openai/private-embedding" }, + model_info: { mode: "embedding" }, + }, + { + model_name: "responses-only", + litellm_params: { model: "openai/private-responses" }, + model_info: { mode: "responses" }, + }, + { model_name: "auto-router", litellm_params: { model: "auto_router/complexity_router" } }, + ]; + vi.mocked(modelInfoCall).mockResolvedValue({ data: deployments, total_pages: 1 }); + const user = userEvent.setup(); + const { start } = mockHooks({}); + await vi.mocked(usePlainModelGroups).withImplementation(hooks.usePlainModelGroups, async () => { + await vi.mocked(usePlainChatModelGroups).withImplementation(hooks.usePlainChatModelGroups, async () => { + render( + + + , + ); + await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "responses-only"); + await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "custom-chat"); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); + await user.click(screen.getByPlaceholderText("Search keys by alias")); + await user.click(within(await screen.findByTestId("paginated-multi-select-list")).getByText("prod-alpha")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); + await user.click(screen.getByPlaceholderText("Select a judge model")); + expect(screen.getAllByRole("option")).toHaveLength(2); + expect(screen.getByRole("option", { name: "custom-chat", exact: true })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "custom-judge", exact: true })).toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "custom-judge", exact: true })); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + expect(screen.getAllByRole("option")).toHaveLength(2); + expect(screen.getByRole("option", { name: "custom-chat", exact: true })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "custom-judge", exact: true })).toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "custom-chat", exact: true })); + await user.click(screen.getByText("Start shadow eval")); + expect(start.mutate).toHaveBeenCalledWith( + expect.objectContaining({ judge_model: "custom-judge", baseline_model: "custom-chat", models: [] }), + ); + }); + }); + client.clear(); + }); + it("offers the start form while the list is still loading", () => { mockHooks({ isPending: true }); render(); @@ -444,7 +532,8 @@ describe("ShadowEvalSection", () => { expect(screen.getByText("Start shadow eval")).toBeDisabled(); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -457,7 +546,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -474,7 +563,7 @@ describe("ShadowEvalSection", () => { await user.click(within(teamList).getByText("engineering")); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -487,7 +576,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -503,7 +592,7 @@ describe("ShadowEvalSection", () => { await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "prod-claude"); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); expect(start.mutate).toHaveBeenCalledWith( @@ -524,20 +613,23 @@ describe("ShadowEvalSection", () => { expect(screen.queryByPlaceholderText("Select a baseline model")).not.toBeInTheDocument(); expect(screen.getByPlaceholderText("Every model the targets use")).toBeInTheDocument(); - await user.click(screen.getByText("Adoption check: key's traffic vs the router")); - await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); expect(screen.queryByPlaceholderText("Every model the targets use")).not.toBeInTheDocument(); await user.click(screen.getByPlaceholderText("Search keys by alias")); const keyList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(keyList).getByText("prod-alpha")); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); expect(screen.getByText("Start shadow eval")).toBeDisabled(); await user.click(screen.getByPlaceholderText("Select a baseline model")); - expect(await screen.findByRole("option", { name: /openai\/gpt-4o/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); await user.click(screen.getByRole("option", { name: /prod-claude/ })); await user.click(screen.getByText("Start shadow eval")); @@ -552,7 +644,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -574,7 +666,7 @@ describe("ShadowEvalSection", () => { screen.getByText("Every router sees the same sampled requests, judged against the same live responses"), ).toBeInTheDocument(); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -587,7 +679,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -605,10 +697,13 @@ describe("ShadowEvalSection", () => { await user.click(await screen.findByText("gpt-auto")); await user.click(routerInput); await user.click(await screen.findByText("claude-auto")); - await user.click(screen.getByText("Adoption check: key's traffic vs the router")); - await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByPlaceholderText("Select a baseline model")); await user.click(screen.getByRole("option", { name: /prod-claude/ })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx index 2eb5fa9c945..d85d26a21a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -5,8 +5,13 @@ import React, { useMemo, useState } from "react"; import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { + useAutoRouters, + usePlainChatModelDeployments, + usePlainChatModelGroups, + usePlainModelGroups, +} from "@/app/(dashboard)/hooks/models/useModels"; +import { buildModelAvailability, deploymentRefsFromModelInfo, resolveAvailableModels } from "@/lib/autorouter_presets"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; @@ -24,53 +29,8 @@ type ShadowEvalDirection = ShadowEvalJob["direction"]; const MAX_ROUTERS = 4; const MAX_MODELS = 100; - const RECOMMENDED_JUDGE_MODELS = ["anthropic/claude-sonnet-5", "openai/gpt-4o", "gemini/gemini-2.5-pro"] as const; -interface CostMapEntry { - litellm_provider?: string; - mode?: string; -} - -const useChatModelNames = (): string[] => { - const { data: costMap } = useModelCostMap(); - return useMemo(() => { - if (!costMap) return []; - const chatModels = Object.entries(costMap as Record) - .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) - .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); - return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); - }, [costMap]); -}; - -const useJudgeModelOptions = (): SearchSelectOption[] => { - const chatModels = useChatModelNames(); - return useMemo(() => { - const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ - label: model, - value: model, - sublabel: "Recommended", - })); - const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); - const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); - return [...pinned, ...rest]; - }, [chatModels]); -}; - -const useBaselineModelOptions = (): SearchSelectOption[] => { - const configuredGroups = usePlainModelGroups(); - const chatModels = useChatModelNames(); - return useMemo(() => { - const configured = [...configuredGroups] - .toSorted((a, b) => a.localeCompare(b)) - .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); - const rest = chatModels - .filter((model) => !configuredGroups.has(model)) - .map((model) => ({ label: model, value: model })); - return [...configured, ...rest]; - }, [configuredGroups, chatModels]); -}; - const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ { value: "forward", label: "Adoption check: key's traffic vs the router" }, { value: "reverse", label: "Regression check: router's picks vs a baseline" }, @@ -276,13 +236,32 @@ export const StartForm: React.FC = () => { const [judgeModel, setJudgeModel] = useState(""); const [maxBudget, setMaxBudget] = useState("10"); const { data: autoRouters } = useAutoRouters(); - const judgeModelOptions = useJudgeModelOptions(); - const baselineModelOptions = useBaselineModelOptions(); const configuredGroups = usePlainModelGroups(); + const chatGroups = usePlainChatModelGroups(); + const chatDeployments = usePlainChatModelDeployments(); const modelOptions = useMemo( () => [...configuredGroups].toSorted((a, b) => a.localeCompare(b)).map((name) => ({ label: name, value: name })), [configuredGroups], ); + const chatOptions = useMemo( + () => modelOptions.filter((option) => chatGroups.has(option.value)), + [modelOptions, chatGroups], + ); + const chatAvailability = useMemo( + () => buildModelAvailability(chatGroups, deploymentRefsFromModelInfo(chatDeployments)), + [chatDeployments, chatGroups], + ); + const recommendedJudgeModels = useMemo( + () => new Set(RECOMMENDED_JUDGE_MODELS.flatMap((model) => resolveAvailableModels(model, chatAvailability))), + [chatAvailability], + ); + const judgeOptions = useMemo( + () => + chatOptions.map((option) => + recommendedJudgeModels.has(option.value) ? { ...option, sublabel: "Recommended" } : option, + ), + [chatOptions, recommendedJudgeModels], + ); const start = useStartShadowEval(); const routerOptions = useMemo(() => { @@ -434,7 +413,7 @@ export const StartForm: React.FC = () => { {direction === "reverse" && ( { )} { }); }); -describe("selectPlainModelGroups", () => { - it("keeps only non-auto-router model groups", () => { - const deployments: AutoRouterCandidateDeployment[] = [ - { model_name: "smart-router", litellm_params: { model: "auto_router/complexity_router" } }, - { model_name: "claude-haiku", litellm_params: { model: "anthropic/claude-haiku-4-5" } }, - { model_name: "claude-sonnet", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, - { model_name: "cheap-router", litellm_params: { model: "auto_router/adaptive_router" } }, +describe("selectPlainChatModelGroups", () => { + it("keeps chat-capable groups when mode metadata is absent or any sibling is compatible", () => { + const deployments: AutoRouterDeployment[] = [ + { model_name: "no-info" }, + { model_name: "null-info", model_info: null }, + { model_name: "empty-info", model_info: {} }, + { model_name: "missing-mode", model_info: { db_model: false } }, + { model_name: "null-mode", model_info: { mode: null } }, + { model_name: "empty-mode", model_info: { mode: "" } }, + { model_name: "chat", model_info: { mode: "chat", db_model: true } }, + { model_name: "completion", model_info: { mode: "completion" } }, + { model_name: "chat-and-missing", model_info: { mode: "chat" } }, + { model_name: "chat-and-missing" }, + { model_name: "chat-then-embedding", model_info: { mode: "chat" } }, + { model_name: "chat-then-embedding", model_info: { mode: "embedding" } }, + { model_name: "embedding-then-chat", model_info: { mode: "embedding" } }, + { model_name: "embedding-then-chat", model_info: { mode: "chat" } }, + { model_name: "embedding-only", model_info: { mode: "embedding" } }, + { model_name: "speech-only", model_info: { mode: "speech" } }, + { model_name: "shared-router", litellm_params: { model: "openai/gpt-4o" } }, + { model_name: "shared-router", litellm_params: { model: "auto_router/complexity_router" } }, + { model_name: "", model_info: { mode: "chat" } }, ]; - expect(selectPlainModelGroups(deployments)).toEqual(new Set(["claude-haiku", "claude-sonnet"])); - }); - - it("drops a group name that also fronts an auto-router deployment", () => { - const deployments: AutoRouterCandidateDeployment[] = [ - { model_name: "shared-name", litellm_params: { model: "auto_router/complexity_router" } }, - { model_name: "shared-name", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, - ]; - - expect(selectPlainModelGroups(deployments)).toEqual(new Set()); - }); - - it("drops deployments that have no public model_name", () => { - expect(selectPlainModelGroups([{ model_name: "", litellm_params: { model: "openai/gpt-4o" } }])).toEqual(new Set()); + expect(selectPlainChatModelGroups(deployments)).toEqual( + new Set([ + "no-info", + "null-info", + "empty-info", + "missing-mode", + "null-mode", + "empty-mode", + "chat", + "completion", + "chat-and-missing", + "chat-then-embedding", + "embedding-then-chat", + ]), + ); }); }); @@ -1103,6 +1121,47 @@ describe("useAutoRouterModelGroups", () => { expect(modelInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", 3, 1000); }); + it("uses every page for configured chat groups and keeps custom deployments without mode metadata", async () => { + (modelInfoCall as any).mockImplementation((_t: string, _u: string, _r: string, page: number) => + Promise.resolve( + page === 1 + ? { + data: [ + { model_name: "configured-chat", model_info: { mode: "chat" } }, + { model_name: "embedding-only", model_info: { mode: "embedding" } }, + ], + total_pages: 2, + } + : { + data: [ + { model_name: "custom-no-mode", model_info: { db_model: true } }, + { model_name: "speech-only", model_info: { mode: "speech" } }, + ], + total_pages: 2, + }, + ), + ); + + const { result } = renderHook(() => usePlainChatModelGroups(), { wrapper }); + + await waitFor(() => expect(result.current.size).toBe(2)); + expect(result.current).toEqual(new Set(["configured-chat", "custom-no-mode"])); + expect(modelInfoCall).toHaveBeenCalledTimes(2); + }); + + it("returns an empty chat group set while loading and after failure", async () => { + (modelInfoCall as any).mockReturnValueOnce(new Promise(() => {})); + const loading = renderHook(() => usePlainChatModelGroups(), { wrapper }); + expect(loading.result.current).toEqual(new Set()); + loading.unmount(); + + queryClient.clear(); + (modelInfoCall as any).mockRejectedValueOnce(new Error("boom")); + const failed = renderHook(() => usePlainChatModelGroups(), { wrapper }); + await waitFor(() => expect(modelInfoCall).toHaveBeenCalledTimes(2)); + expect(failed.result.current).toEqual(new Set()); + }); + it("returns an empty set before the model list resolves", () => { (modelInfoCall as any).mockReturnValue(new Promise(() => {})); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index b3a783a71dc..579ee7ff81a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -2,6 +2,7 @@ import { useQuery, useInfiniteQuery, useQueryClient, UseQueryResult } from "@tan import { createQueryKeys } from "../common/queryKeysFactory"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; +import { EndpointType, isModeCompatibleWithEndpoint } from "@/components/chat_ui/mode_endpoint_mapping"; export interface ProxyModel { id: string; @@ -87,6 +88,7 @@ export const useModelsInfo = ( const AUTO_ROUTER_MODEL_PREFIX = "auto_router/"; const AUTO_ROUTER_LOOKUP_PAGE_SIZE = 1000; const NO_AUTO_ROUTERS: ReadonlySet = new Set(); +const NO_DEPLOYMENTS: AutoRouterDeployment[] = []; export interface AutoRouterCandidateDeployment { model_name?: string | null; @@ -96,6 +98,7 @@ export interface AutoRouterCandidateDeployment { export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { litellm_params?: { model?: string | null; + base_model?: string | null; complexity_router_config?: unknown; complexity_router_default_model?: string | null; auto_router_config?: unknown; @@ -111,6 +114,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { /** False for config.yaml-defined deployments, which the update and delete routes refuse. */ db_model?: boolean | null; base_model?: string | null; + mode?: string | null; created_at?: string | null; updated_at?: string | null; team_id?: string | null; @@ -142,6 +146,22 @@ export const selectPlainModelGroups = (deployments: AutoRouterCandidateDeploymen ); }; +export const selectPlainChatModelDeployments = (deployments: AutoRouterDeployment[]): AutoRouterDeployment[] => { + const plainGroups = selectPlainModelGroups(deployments); + return deployments.filter( + (deployment) => + plainGroups.has(deployment.model_name ?? "") && + isModeCompatibleWithEndpoint(deployment.model_info?.mode, EndpointType.CHAT), + ); +}; + +export const selectPlainChatModelGroups = (deployments: AutoRouterDeployment[]): ReadonlySet => + new Set( + selectPlainChatModelDeployments(deployments) + .map((deployment) => deployment.model_name) + .filter((name): name is string => Boolean(name)), + ); + export const fetchAllModelDeployments = async ( accessToken: string, userId: string, @@ -180,37 +200,32 @@ export const autoRouterListKey = (userId: string | null, userRole: string | null }, }); -export const useAutoRouterModelGroups = (): ReadonlySet => { +const useDeployments = ( + select: (deployments: AutoRouterDeployment[]) => TSelected, +): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); - const { data } = useQuery>({ + return useQuery({ queryKey: autoRouterListKey(userId, userRole), queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), enabled: Boolean(accessToken && userId && userRole), - select: selectAutoRouterModelGroups, + select, }); - return data ?? NO_AUTO_ROUTERS; }; -export const usePlainModelGroups = (): ReadonlySet => { - const { accessToken, userId, userRole } = useAuthorized(); - const { data } = useQuery>({ - queryKey: autoRouterListKey(userId, userRole), - queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), - enabled: Boolean(accessToken && userId && userRole), - select: selectPlainModelGroups, - }); - return data ?? NO_AUTO_ROUTERS; -}; +export const useAutoRouterModelGroups = (): ReadonlySet => + useDeployments(selectAutoRouterModelGroups).data ?? NO_AUTO_ROUTERS; -export const useAutoRouters = (): UseQueryResult => { - const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ - queryKey: autoRouterListKey(userId, userRole), - queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), - enabled: Boolean(accessToken && userId && userRole), - select: selectAutoRouterDeployments, - }); -}; +export const usePlainModelGroups = (): ReadonlySet => + useDeployments(selectPlainModelGroups).data ?? NO_AUTO_ROUTERS; + +export const usePlainChatModelGroups = (): ReadonlySet => + useDeployments(selectPlainChatModelGroups).data ?? NO_AUTO_ROUTERS; + +export const usePlainChatModelDeployments = (): AutoRouterDeployment[] => + useDeployments(selectPlainChatModelDeployments).data ?? NO_DEPLOYMENTS; + +export const useAutoRouters = (): UseQueryResult => + useDeployments(selectAutoRouterDeployments); export const useInvalidateAutoRouters = (): (() => Promise) => { const queryClient = useQueryClient(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx index 8fd4a57dbfd..87d569f6490 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx @@ -1,7 +1,9 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { EndpointType, getEndpointType, ModelMode } from "@/components/chat_ui/mode_endpoint_mapping"; - -const KNOWN_MODEL_MODES = new Set(Object.values(ModelMode)); +import { + EndpointType, + getEndpointType, + isModeCompatibleWithEndpoint, +} from "@/components/chat_ui/mode_endpoint_mapping"; export const determineEndpointType = (selectedModel: string, modelInfo: ModelGroup[]): EndpointType => { const selectedModelInfo = modelInfo.find((option) => option.model_group === selectedModel); @@ -13,31 +15,8 @@ export const determineEndpointType = (selectedModel: string, modelInfo: ModelGro return EndpointType.CHAT; }; -export const isModelCompatibleWithEndpoint = (model: ModelGroup, endpointType: EndpointType): boolean => { - if (!model.mode) { - return true; - } - - if (!KNOWN_MODEL_MODES.has(model.mode)) { - return false; - } - - const optionEndpoint = getEndpointType(model.mode); - - if ( - endpointType === EndpointType.RESPONSES || - endpointType === EndpointType.ANTHROPIC_MESSAGES || - endpointType === EndpointType.INTERACTIONS - ) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; - } - - if (endpointType === EndpointType.IMAGE_EDITS) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; - } - - return optionEndpoint === endpointType; -}; +export const isModelCompatibleWithEndpoint = (model: ModelGroup, endpointType: EndpointType): boolean => + isModeCompatibleWithEndpoint(model.mode, endpointType); export const filterModelsForEndpoint = (models: ModelGroup[], endpointType: EndpointType): ModelGroup[] => models.filter((model) => isModelCompatibleWithEndpoint(model, endpointType)); diff --git a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx index 930ded5d1a5..b46200ccbf2 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx @@ -57,3 +57,20 @@ export const getEndpointType = (mode: string): EndpointType => { // else default to chat return EndpointType.CHAT; }; + +export const isModeCompatibleWithEndpoint = (mode: string | null | undefined, endpointType: EndpointType): boolean => { + if (!mode) return true; + if (!Object.values(ModelMode).includes(mode as ModelMode)) return false; + const optionEndpoint = getEndpointType(mode); + if ( + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES || + endpointType === EndpointType.INTERACTIONS + ) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; + } + if (endpointType === EndpointType.IMAGE_EDITS) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; + } + return optionEndpoint === endpointType; +}; diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 8a5f83adbdf..fed11454c23 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -13,6 +13,7 @@ import { buildModelAvailability, deploymentRefsFromModelInfo, normalizeModelName, + resolveAvailableModels, } from "./autorouter_presets"; import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching"; import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords"; @@ -380,6 +381,18 @@ describe("autorouter_presets", () => { expect(availability.underlyingIndex.size).toBe(0); }); + it("returns every configured group serving the same underlying model", () => { + const availability = buildModelAvailability( + ["z-group", "a-group"], + [ + { modelGroup: "z-group", underlyingModels: ["anthropic/claude-sonnet-5"] }, + { modelGroup: "a-group", underlyingModels: ["bedrock/us.anthropic.claude-sonnet-5-v1:0"] }, + ], + ); + + expect(resolveAvailableModels("anthropic/claude-sonnet-5", availability)).toEqual(["a-group", "z-group"]); + }); + it("breaks ties between groups serving the same model deterministically, alphabetically", () => { const availability = buildModelAvailability( ["z-group", "a-group"], diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index f2df55fc310..02096cada41 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -159,16 +159,19 @@ export const deploymentRefsFromModelInfo = ( return row.model_name && underlyingModels.length > 0 ? [{ modelGroup: row.model_name, underlyingModels }] : []; }); -export const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => { +export const resolveAvailableModels = (requiredModel: string, availability: ModelAvailability): readonly string[] => { const { modelGroups, underlyingIndex } = availability; - if (modelGroups.has(requiredModel)) return requiredModel; + if (modelGroups.has(requiredModel)) return [requiredModel]; const normalized = normalizeModelName(requiredModel); - const groupMatch = Array.from(modelGroups).find((available) => normalizeModelName(available) === normalized); - if (groupMatch !== undefined) return groupMatch; + const groupMatches = Array.from(modelGroups).filter((available) => normalizeModelName(available) === normalized); + if (groupMatches.length > 0) return groupMatches; const key = normalizeUnderlyingModel(requiredModel); - return key === null ? undefined : underlyingIndex.get(key)?.[0]; + return key === null ? [] : underlyingIndex.get(key) ?? []; }; +export const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => + resolveAvailableModels(requiredModel, availability)[0]; + export const getMissingModels = ( config: Parameters[0], availability: ModelAvailability, From 3aa9267df278a4a28399460cc7cbe4708a4bd169 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 18:19:52 -0700 Subject: [PATCH 156/157] fix(ui): drop blank streamed costs instead of reading them as zero Number("") and Number(" ") both return 0, which passed the finite check, so a provider reporting an empty cost got a fabricated $0.000000 metric instead of having the unusable value omitted. Both ingestion sites carried the same inline parsing, so this pulls it into one parseUsageCost helper that keeps finite numbers and non-blank numeric strings and drops everything else, including booleans, arrays and breakdown objects. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YNw8WvkvCSeTcE5qvergu3 --- .../llm_calls/chat_completion.test.tsx | 6 +++ .../components/llm_calls/chat_completion.tsx | 10 ++--- .../llm_calls/responses_api.test.tsx | 29 ++++++++++++++ .../components/llm_calls/responses_api.tsx | 9 ++--- .../components/llm_calls/usage_cost.test.ts | 39 +++++++++++++++++++ .../src/components/llm_calls/usage_cost.ts | 23 +++++++++++ 6 files changed, 105 insertions(+), 11 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts create mode 100644 ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx index 129d51d50a2..71f36bdb6db 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx @@ -472,6 +472,12 @@ describe("chat_completion prompt cache usage", () => { expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); }); + + it("omits cost when the provider reports a blank value", async () => { + const usageData = await captureUsage({ cost: " " }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); }); describe("chat_completion response cache", () => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index c12a2c3cfb4..ffa2877fbd9 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -5,6 +5,7 @@ import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk => ({ @@ -243,12 +244,9 @@ export async function makeOpenAIChatCompletionRequest( usageData.reasoningTokens = chunkWithUsage.usage.completion_tokens_details.reasoning_tokens; } - // Extract cost from usage object if available - if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - const parsedCost = Number(chunkWithUsage.usage.cost); - if (Number.isFinite(parsedCost)) { - usageData.cost = parsedCost; - } + const parsedCost = parseUsageCost(chunkWithUsage.usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData); diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 7b13f0ac2e3..b55df89d5cc 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -260,6 +260,35 @@ describe("responses_api", () => { expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); }); + it("should omit cost when the proxy reports a blank cost", async () => { + async function* streamWithBlankCost() { + yield { + type: "response.completed", + response: { + id: "resp_blank_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: " " }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithBlankCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + it("should replay MCP output items as events for a non-streaming response", async () => { mockResponsesCreate.mockReturnValueOnce( nonStreamingResponse({ diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 63085a81b94..ce54c7c6b40 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; import type { MCPEvent } from "@/components/mcp_tools/types"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { @@ -311,11 +312,9 @@ export async function makeOpenAIResponsesRequest( usageData.reasoningTokens = reasoningTokens; } - if (usage.cost !== undefined && usage.cost !== null) { - const parsedCost = Number(usage.cost); - if (Number.isFinite(parsedCost)) { - usageData.cost = parsedCost; - } + const parsedCost = parseUsageCost(usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData, mcpToolUsed); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts new file mode 100644 index 00000000000..ee1021c0629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from "vitest"; +import { parseUsageCost } from "./usage_cost"; + +describe("parseUsageCost", () => { + it("keeps finite numbers, including zero", () => { + expect(parseUsageCost(0)).toBe(0); + expect(parseUsageCost(0.000063)).toBe(0.000063); + }); + + it("keeps numeric strings", () => { + expect(parseUsageCost("0.00019")).toBe(0.00019); + expect(parseUsageCost(" 0.00019 ")).toBe(0.00019); + }); + + it("drops blank strings instead of fabricating a zero cost", () => { + expect(parseUsageCost("")).toBeUndefined(); + expect(parseUsageCost(" ")).toBeUndefined(); + expect(parseUsageCost("\t\n")).toBeUndefined(); + }); + + it("drops strings with a numeric prefix instead of truncating them", () => { + expect(parseUsageCost("1oops")).toBeUndefined(); + expect(parseUsageCost("0.5 USD")).toBeUndefined(); + }); + + it("drops non-finite numbers", () => { + expect(parseUsageCost(Number.NaN)).toBeUndefined(); + expect(parseUsageCost(Number.POSITIVE_INFINITY)).toBeUndefined(); + }); + + it("drops values that are not numbers or strings", () => { + expect(parseUsageCost(null)).toBeUndefined(); + expect(parseUsageCost(undefined)).toBeUndefined(); + expect(parseUsageCost(true)).toBeUndefined(); + expect(parseUsageCost([])).toBeUndefined(); + expect(parseUsageCost(["0.5"])).toBeUndefined(); + expect(parseUsageCost({ total_cost: 0.5 })).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts new file mode 100644 index 00000000000..79f56dcb2df --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts @@ -0,0 +1,23 @@ +/** + * Providers and upstream gateways report `usage.cost` unvalidated: it arrives as a number, a numeric + * string, an empty string, or something non-numeric. A cost that does not resolve to a finite number + * must be dropped rather than coerced, because `NaN` survives `JSON.stringify` as `null` and crashes + * the metrics row on the next load. + */ +export function parseUsageCost(rawCost: unknown): number | undefined { + if (typeof rawCost === "number") { + return Number.isFinite(rawCost) ? rawCost : undefined; + } + + if (typeof rawCost !== "string") { + return undefined; + } + + const trimmed = rawCost.trim(); + if (trimmed === "") { + return undefined; + } + + const parsed = Number(trimmed); + return Number.isFinite(parsed) ? parsed : undefined; +} From 4adf99557f901fd73525697fdba8c9b2c99d0cc9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:27:12 -0700 Subject: [PATCH 157/157] fix(responses): echo "auto" for a tool_choice the bridge cannot express instead of failing after the provider call --- .../transformation.py | 6 +++-- .../test_litellm_completion_responses.py | 26 +++++++++++++++++++ .../test_streaming_iterator_transformation.py | 20 ++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 6b9931b7df2..fca5b0d11cf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -287,8 +287,10 @@ class LiteLLMCompletionResponsesConfig: return ToolChoiceCustomParam(type="custom", name=custom_name) case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _, normalized: - return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(normalized) + case _, "none" | "auto" | "required" as normalized: + return normalized + case _, _: + return "auto" @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 8f2937cb252..46249e50572 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1435,6 +1435,9 @@ class TestToolChoiceTransformation: ("required", "required"), ("none", "none"), (None, "auto"), + ("any", "auto"), + ("run_command", "auto"), + ({"name": "run_command"}, "auto"), ], ) def test_transform_tool_choice_for_responses_api_response( @@ -1478,6 +1481,29 @@ class TestToolChoiceTransformation: assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + def test_non_streamed_response_with_unrecognized_tool_choice_echoes_auto(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-unrecognized-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message(role="assistant", content="/Users/dev"), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": "any"}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == "auto" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 70f0957890d..850ee7ba623 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -685,3 +685,23 @@ def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: {"type": "function", "name": "run_command"}, ] assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) + + +def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": "any", + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + response_events: Final = [ + event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"]