From 518c4b07a72e989ed777f93c01ef296c77a8b567 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 21:55:44 +0000 Subject: [PATCH 01/89] fix(azure_ai): route Responses API to native /openai/v1/responses for Foundry Models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/llms/azure_ai/chat/transformation.py | 8 +- litellm/llms/azure_ai/common_utils.py | 26 +++ litellm/llms/azure_ai/responses/__init__.py | 0 .../llms/azure_ai/responses/transformation.py | 62 +++++++ litellm/utils.py | 8 + .../llms/azure_ai/responses/__init__.py | 0 .../test_azure_ai_responses_transformation.py | 173 ++++++++++++++++++ 9 files changed, 279 insertions(+), 6 deletions(-) create mode 100644 litellm/llms/azure_ai/responses/__init__.py create mode 100644 litellm/llms/azure_ai/responses/transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/responses/__init__.py create mode 100644 tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 2f6643c644c..6daf18e43ad 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1742,6 +1742,9 @@ if TYPE_CHECKING: from .llms.azure.responses.o_series_transformation import ( AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig, ) + from .llms.azure_ai.responses.transformation import ( + AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig, + ) from .llms.xai.responses.transformation import ( XAIResponsesAPIConfig as XAIResponsesAPIConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 488331e3895..744de95cab6 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -231,6 +231,7 @@ LLM_CONFIG_NAMES = ( "OpenAIResponsesAPIConfig", "AzureOpenAIResponsesAPIConfig", "AzureOpenAIOSeriesResponsesAPIConfig", + "AzureAIResponsesAPIConfig", "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "HostedVLLMResponsesAPIConfig", @@ -935,6 +936,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.azure.responses.o_series_transformation", "AzureOpenAIOSeriesResponsesAPIConfig", ), + "AzureAIResponsesAPIConfig": ( + ".llms.azure_ai.responses.transformation", + "AzureAIResponsesAPIConfig", + ), "XAIResponsesAPIConfig": ( ".llms.xai.responses.transformation", "XAIResponsesAPIConfig", diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 27a98347087..27e4f405dfd 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,7 +1,6 @@ import enum import re from typing import Any, List, Optional, Tuple, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -12,6 +11,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( _audio_or_image_in_message_content, convert_content_list_to_str, ) +from litellm.llms.azure_ai.common_utils import azure_ai_use_api_key_header from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error @@ -85,11 +85,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url = urlparse(api_base) - host = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return azure_ai_use_api_key_header(api_base) def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 9965aa693c3..26021367440 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,4 +1,5 @@ from typing import List, Literal, Optional +from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter @@ -6,6 +7,31 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +def azure_ai_use_api_key_header(api_base: str) -> bool: + """Whether Azure AI auth should use the `api-key` header instead of a Bearer token. + + Foundry and Azure OpenAI hosts authenticate key-based requests with the + `api-key` header; serverless/other endpoints expect `Authorization: Bearer`. + """ + host = urlparse(api_base).hostname + return bool(host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com"))) + + +def azure_ai_supports_native_responses(model: str | None) -> bool: + """Whether an Azure AI model should use the native Responses API rather than the chat bridge. + + Foundry Models expose an OpenAI-compatible Responses endpoint at + `/openai/v1/responses`. Claude deployments speak the Anthropic + Messages API and the model-router/agents routes have their own surfaces, so + those keep the chat-completions bridge. + """ + if not model: + return False + if "claude" in model.lower(): + return False + return AzureFoundryModelInfo.get_azure_ai_route(model) == "default" + + class AzureFoundryModelInfo(BaseLLMModelInfo): """Model info for Azure AI / Azure Foundry models.""" diff --git a/litellm/llms/azure_ai/responses/__init__.py b/litellm/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/azure_ai/responses/transformation.py b/litellm/llms/azure_ai/responses/transformation.py new file mode 100644 index 00000000000..43fb9fc93d0 --- /dev/null +++ b/litellm/llms/azure_ai/responses/transformation.py @@ -0,0 +1,62 @@ +import httpx + +from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + azure_ai_use_api_key_header, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import _add_path_to_api_base + + +class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): + """Native Responses API config for Azure AI Foundry Models. + + Foundry Models such as the GPT-5 family expose an OpenAI-compatible Responses + endpoint at `/openai/v1/responses`. Routing here (instead of the + chat-completions bridge) keeps `reasoning_effort` alongside function tools, + which Azure rejects on `/chat/completions`. + """ + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.AZURE_AI + + def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: + litellm_params = litellm_params or GenericLiteLLMParams() + api_key = AzureFoundryModelInfo.get_api_key(litellm_params.api_key) + api_base = AzureFoundryModelInfo.get_api_base(litellm_params.api_base) + + if api_key: + if api_base and azure_ai_use_api_key_header(api_base): + headers["api-key"] = api_key + else: + headers["Authorization"] = f"Bearer {api_key}" + else: + headers = BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) + + headers.setdefault("Content-Type", "application/json") + return headers + + def get_complete_url( + self, + api_base: str | None, + litellm_params: dict, + ) -> str: + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "api_base is required for Azure AI Foundry Responses API. " + "Set the api_base parameter or the AZURE_AI_API_BASE environment variable." + ) + + original_url = httpx.URL(api_base) + query_params = dict(original_url.params) + api_version = litellm_params.get("api_version") + if "api-version" not in query_params and isinstance(api_version, str): + query_params["api-version"] = api_version + + new_url = _add_path_to_api_base(api_base=api_base, ending_path="/openai/v1/responses") + return str(httpx.URL(new_url).copy_with(params=query_params)) diff --git a/litellm/utils.py b/litellm/utils.py index 174bed09396..34a083786d9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8221,6 +8221,14 @@ class ProviderConfigManager: return litellm.AzureOpenAIOSeriesResponsesAPIConfig() else: return litellm.AzureOpenAIResponsesAPIConfig() + elif litellm.LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.common_utils import ( + azure_ai_supports_native_responses, + ) + + if azure_ai_supports_native_responses(model): + return litellm.AzureAIResponsesAPIConfig() + return None elif litellm.LlmProviders.XAI == provider: return litellm.XAIResponsesAPIConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: diff --git a/tests/test_litellm/llms/azure_ai/responses/__init__.py b/tests/test_litellm/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py new file mode 100644 index 00000000000..0a2f2a4ab8e --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -0,0 +1,173 @@ +""" +Regression tests for native Azure AI Foundry Responses API routing (LIT-4427). + +Before the fix, `azure_ai` had no native Responses config, so `litellm.responses()` +fell back to the chat-completions bridge and sent `reasoning_effort` + function tools +to `/chat/completions`, which Azure rejects for GPT-5 models. These tests assert the +request now goes to the native `/openai/v1/responses` endpoint in Responses shape. +""" + +import json +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +import litellm +from litellm.llms.azure_ai.responses.transformation import AzureAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager + + +class MockResponse: + def __init__(self, json_data, status_code=200): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + self.headers = httpx.Headers({}) + + def json(self): + return self._json_data + + +def _minimal_responses_payload(model: str) -> dict: + return { + "id": "resp_123", + "object": "response", + "created_at": 1741369938, + "status": "completed", + "model": model, + "output": [], + "parallel_tool_calls": False, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "error": None, + "tool_choice": "auto", + "tools": [], + "metadata": None, + "temperature": None, + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "instructions": None, + "incomplete_details": None, + "user": None, + } + + +@pytest.mark.parametrize( + "model", + ["gpt-5.6-luna-20260710154139", "gpt-5.5-20260504143601", "DeepSeek-R1-0528"], +) +def test_azure_ai_resolves_native_responses_config(model): + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model) + assert isinstance(config, AzureAIResponsesAPIConfig) + + +@pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"]) +def test_azure_ai_non_responses_models_keep_bridge(model): + """Claude / model-router / agents routes have their own surfaces, so they must + keep returning None (chat-completions bridge).""" + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model) + assert config is None + + +@pytest.mark.parametrize( + "api_base,expected", + [ + ( + "https://res.services.ai.azure.com/api/projects/proj", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com/api/projects/proj/", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com", + "https://res.services.ai.azure.com/openai/v1/responses", + ), + ( + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ], +) +def test_get_complete_url(api_base, expected): + config = AzureAIResponsesAPIConfig() + assert config.get_complete_url(api_base=api_base, litellm_params={}) == expected + + +def test_validate_environment_api_key_header_for_foundry_host(): + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams( + api_key="secret", api_base="https://res.services.ai.azure.com/api/projects/proj" + ), + ) + assert headers["api-key"] == "secret" + assert "Authorization" not in headers + + +def test_validate_environment_bearer_for_serverless_host(): + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams( + api_key="secret", api_base="https://endpoint.eastus.models.ai.azure.com" + ), + ) + assert headers["Authorization"] == "Bearer secret" + assert "api-key" not in headers + + +@pytest.mark.asyncio +async def test_aresponses_routes_to_native_endpoint_with_reasoning_and_tools(): + """Core LIT-4427 regression: reasoning_effort + function tools must be sent to the + native /openai/v1/responses endpoint in Responses shape, not bridged to /chat/completions.""" + tools = [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ] + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(_minimal_responses_payload("gpt-5.6-luna"), 200) + + await litellm.aresponses( + model="azure_ai/gpt-5.6-luna-20260710154139", + input="What is the weather in SF?", + reasoning_effort="high", + tools=tools, + api_base="https://res.services.ai.azure.com/api/projects/proj", + api_key="fake-key", + ) + + mock_post.assert_called_once() + url = str(mock_post.call_args.kwargs["url"]) + body = mock_post.call_args.kwargs["json"] + + assert url == "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses" + assert "/chat/completions" not in url + assert "input" in body + assert "messages" not in body + assert body["reasoning"] == {"effort": "high"} + assert body["tools"] == tools From 34d32c04e5cdf6a7ce88e8b3f359fa1b4d891d73 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 22:06:58 +0000 Subject: [PATCH 02/89] test(azure_ai): drop responses test __init__ to fix package name collision; cover api-version, missing api_base, AD fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/azure_ai/responses/__init__.py | 0 .../test_azure_ai_responses_transformation.py | 35 +++++++++++++++++++ 2 files changed, 35 insertions(+) delete mode 100644 tests/test_litellm/llms/azure_ai/responses/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/responses/__init__.py b/tests/test_litellm/llms/azure_ai/responses/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 0a2f2a4ab8e..b853b114645 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -103,6 +103,25 @@ def test_get_complete_url(api_base, expected): assert config.get_complete_url(api_base=api_base, litellm_params={}) == expected +def test_get_complete_url_adds_api_version_from_params(): + config = AzureAIResponsesAPIConfig() + url = config.get_complete_url( + api_base="https://res.services.ai.azure.com/api/projects/proj", + litellm_params={"api_version": "2025-04-01-preview"}, + ) + assert url == ( + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses?api-version=2025-04-01-preview" + ) + + +def test_get_complete_url_raises_without_api_base(monkeypatch): + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + config = AzureAIResponsesAPIConfig() + with pytest.raises(ValueError): + config.get_complete_url(api_base=None, litellm_params={}) + + def test_validate_environment_api_key_header_for_foundry_host(): config = AzureAIResponsesAPIConfig() headers = config.validate_environment( @@ -129,6 +148,22 @@ def test_validate_environment_bearer_for_serverless_host(): assert "api-key" not in headers +def test_validate_environment_falls_back_to_base_azure_env_without_key(monkeypatch): + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams(api_base="https://res.services.ai.azure.com/api/projects/proj"), + ) + assert headers["Content-Type"] == "application/json" + assert "api-key" not in headers + + @pytest.mark.asyncio async def test_aresponses_routes_to_native_endpoint_with_reasoning_and_tools(): """Core LIT-4427 regression: reasoning_effort + function tools must be sent to the From 538a86885d11323a2883088947362d6b50f2e8fd Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 22:17:03 +0000 Subject: [PATCH 03/89] fix(azure_ai): return native Responses config for management ops (model=None) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 2 +- .../responses/test_azure_ai_responses_transformation.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 34a083786d9..20d5993e14c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8226,7 +8226,7 @@ class ProviderConfigManager: azure_ai_supports_native_responses, ) - if azure_ai_supports_native_responses(model): + if model is None or azure_ai_supports_native_responses(model): return litellm.AzureAIResponsesAPIConfig() return None elif litellm.LlmProviders.XAI == provider: diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index b853b114645..1d59773728a 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -65,6 +65,14 @@ def test_azure_ai_resolves_native_responses_config(model): assert isinstance(config, AzureAIResponsesAPIConfig) +def test_azure_ai_resolves_native_config_for_management_ops(): + """Management ops (delete/get/cancel/list) call the lookup with model=None; it must + still return the native config so those operations can build the right URL after a + native create succeeds.""" + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=None) + assert isinstance(config, AzureAIResponsesAPIConfig) + + @pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"]) def test_azure_ai_non_responses_models_keep_bridge(model): """Claude / model-router / agents routes have their own surfaces, so they must From 62d8258868b745bfab8f7521ba498ad2a1564b06 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:27:50 +0000 Subject: [PATCH 04/89] fix(dashscope): forward reasoning_effort to the provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/dashscope/chat/transformation.py | 6 ++++++ .../test_dashscope_chat_transformation.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 5ab7fbf3658..977bb38f59a 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class DashScopeChatConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a list + return [ # mutable-ok: base class contract returns a list + *super().get_supported_openai_params(model=model), + "reasoning_effort", + ] + def remove_cache_control_flag_from_messages_and_tools( self, model: str, diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py index 8dbc197d4b5..e3bbf2abc48 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py @@ -169,6 +169,21 @@ class TestDashScopeConfig: assert transformed_messages[0].get("cache_control") == {"type": "ephemeral"} + @pytest.mark.parametrize("reasoning_effort", ["none", "minimal", "low", "high"]) + def test_dashscope_forwards_reasoning_effort(self, reasoning_effort: str): + """DashScope supports reasoning_effort, so it must reach the provider instead of being dropped.""" + assert "reasoning_effort" in DashScopeChatConfig().get_supported_openai_params( + model="qwen3.7-plus" + ) + + optional_params = litellm.get_optional_params( + model="qwen3.7-plus", + custom_llm_provider="dashscope", + reasoning_effort=reasoning_effort, + ) + + assert optional_params["reasoning_effort"] == reasoning_effort + def test_dashscope_preserves_cache_control_in_tools(self): """DashScope should NOT strip cache_control from tools.""" config = DashScopeChatConfig() From 2b6c30c2d93199586fd696fab7cedc0c5f769b16 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:42:38 +0000 Subject: [PATCH 05/89] refactor(dashscope): tighten supported params return type Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/dashscope/chat/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 977bb38f59a..04d530ea89b 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -12,7 +12,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class DashScopeChatConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a list + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list return [ # mutable-ok: base class contract returns a list *super().get_supported_openai_params(model=model), "reasoning_effort", From dff08dcb55b35ff9346445b0ba3a733692e249f1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:16:45 +0000 Subject: [PATCH 06/89] fix(proxy): retry rate-limit fallbacks from a pristine request snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 38 ++- .../proxy/test_common_request_processing.py | 240 ++++++++++++++++++ 2 files changed, 264 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e3a2b892721..255bf3ebda2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -32,7 +32,11 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error +from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + independent_snapshot, + is_expected_client_error, +) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -2034,6 +2038,21 @@ class ProxyBaseLLMRequestProcessing: ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + original_model: Final = self.data.get("model") + fallback_models: Final = ( + self._resolve_fallback_models( + model=original_model, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + if original_model + and isinstance(original_model, str) + and llm_router + and not self.data.get("disable_fallbacks") + else None + ) + pristine: Final = independent_snapshot(self.data) if fallback_models else None + try: return await self.common_processing_pre_call_logic( request=request, @@ -2052,16 +2071,7 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) except ProxyRateLimitError as original_exc: - original_model: Final = self.data.get("model") - if not original_model or not llm_router or self.data.get("disable_fallbacks"): - raise - - fallback_models: Final = self._resolve_fallback_models( - model=original_model, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - ) - if not fallback_models: + if not fallback_models or pristine is None: raise verbose_proxy_logger.info( @@ -2074,7 +2084,7 @@ class ProxyBaseLLMRequestProcessing: for fallback_model in fallback_models: if fallback_model == original_model: continue - self.data["model"] = fallback_model + self.data = {**independent_snapshot(pristine), "model": fallback_model} try: return await self.common_processing_pre_call_logic( request=request, @@ -2095,10 +2105,10 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError: continue except BaseException: - self.data["model"] = original_model + self.data = pristine raise - self.data["model"] = original_model + self.data = pristine raise original_exc def _resolve_fallback_models( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 50b26577e5c..89e0799ba16 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6235,6 +6235,246 @@ class TestPreCallWithFallbacksOnLocalRateLimit: call_type="acompletion", ) + @pytest.mark.asyncio + async def test_fallback_retries_from_pristine_request_data(self): + import threading + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4" + fallback_model = "gpt-3.5-turbo" + + processor = ProxyBaseLLMRequestProcessing( + data={ + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + ) + + metadata_at_entry = [] + + async def mock_pre_call_logic(**kwargs): + copy.deepcopy(processor.data["metadata"]) + metadata_at_entry.append(dict(processor.data["metadata"])) + processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() + processor.data["litellm_logging_obj"] = object() + if processor.data.get("model") == primary_model: + raise ProxyRateLimitError( + detail="TPM limit exceeded for gpt-4", + headers={"retry-after": "30"}, + ) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: [fallback_model]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + data, logging_obj = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == fallback_model + assert metadata_at_entry[1] == {"tags": ["a"]} + + @pytest.mark.asyncio + async def test_exhausted_fallbacks_restore_pristine_request_data(self): + import threading + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4" + original_data = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + processor = ProxyBaseLLMRequestProcessing(data=copy.deepcopy(original_data)) + + async def mock_pre_call_logic(**kwargs): + processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() + processor.data["litellm_logging_obj"] = object() + raise ProxyRateLimitError( + detail=f"TPM limit exceeded for {processor.data.get('model')}", + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: ["gpt-3.5-turbo"]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError, match="gpt-4"): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data == original_data + + @pytest.mark.asyncio + async def test_real_add_litellm_data_to_request_rerun_with_otel_span_falls_back(self): + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import ProxyConfig + + trace.set_tracer_provider(TracerProvider()) + + primary_model = "gpt-4" + fallback_model = "gpt-3.5-turbo" + + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + user_api_key_dict = UserAPIKeyAuth( + parent_otel_span=trace.get_tracer("x").start_span("s"), + api_key="hashed-key", + user_id="u1", + team_id="t1", + metadata={}, + team_metadata={}, + team_member_tpm_limit=1000, + ) + + processor = ProxyBaseLLMRequestProcessing( + data={ + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + ) + + async def real_add_litellm_data_pre_call(**kwargs): + await add_litellm_data_to_request( + data=processor.data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=ProxyConfig(), + general_settings={}, + version="test", + ) + if processor.data.get("model") == primary_model: + raise ProxyRateLimitError( + detail="TPM limit exceeded for gpt-4", + headers={"retry-after": "30"}, + ) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: [fallback_model]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=real_add_litellm_data_pre_call, + ): + data, logging_obj = await processor._pre_call_with_fallbacks( + request=request_mock, + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=user_api_key_dict, + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == fallback_model + + @pytest.mark.asyncio + async def test_no_fallbacks_skips_snapshot(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + raise ProxyRateLimitError( + detail="TPM limit exceeded", + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = None + + with patch("litellm.proxy.common_request_processing.independent_snapshot") as snapshot_mock: + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + snapshot_mock.assert_not_called() + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From c064e576ee31ed05b067462412c6c234d364c8fd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:29:20 +0000 Subject: [PATCH 07/89] fix(proxy): tolerate missing router_settings and non-list fallbacks in fallback resolution Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 4 ++-- tests/test_litellm/proxy/test_common_request_processing.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 255bf3ebda2..9927320d795 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2121,14 +2121,14 @@ class ProxyBaseLLMRequestProcessing: fallbacks = None - key_router_settings: Final = user_api_key_dict.router_settings + key_router_settings: Final = getattr(user_api_key_dict, "router_settings", None) if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: fallbacks = key_router_settings["fallbacks"] if fallbacks is None: fallbacks = llm_router.fallbacks - if not fallbacks: + if not isinstance(fallbacks, list) or not fallbacks: return None fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 89e0799ba16..4beb73b3a81 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6449,7 +6449,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = None - with patch("litellm.proxy.common_request_processing.independent_snapshot") as snapshot_mock: + with patch( # test-quality-ok: spying the snapshot seam is the only observable check that the no-fallback path skips it + "litellm.proxy.common_request_processing.independent_snapshot" + ) as snapshot_mock: with patch.object( processor, "common_processing_pre_call_logic", From 431dcce6a72ddcb6ba73c56848e1b5e0477b7699 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:46:44 +0000 Subject: [PATCH 08/89] test(proxy): give pre-call mocks real router_settings and fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 4 ++-- .../test_response_polling_pre_call_checks.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9927320d795..255bf3ebda2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2121,14 +2121,14 @@ class ProxyBaseLLMRequestProcessing: fallbacks = None - key_router_settings: Final = getattr(user_api_key_dict, "router_settings", None) + key_router_settings: Final = user_api_key_dict.router_settings if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: fallbacks = key_router_settings["fallbacks"] if fallbacks is None: fallbacks = llm_router.fallbacks - if not isinstance(fallbacks, list) or not fallbacks: + if not fallbacks: return None fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..38f087f51ca 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -48,10 +48,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(), + llm_router=MagicMock(fallbacks=None), general_settings={}, proxy_config=MagicMock(), skip_pre_call_logic=True, @@ -87,10 +87,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(), + llm_router=MagicMock(fallbacks=None), general_settings={}, proxy_config=MagicMock(), ) From 9fbc2c5b71251d879d02cb8269fd0191fe7ec448 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:23:06 +0000 Subject: [PATCH 09/89] fix(proxy): show all model groups to proxy admins in /model_group/info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 32 +++++---- .../proxy_server/test_routes_model_info.py | 70 ++++++++++++++++--- 2 files changed, 79 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..841ac47399e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15494,18 +15494,26 @@ async def model_group_info( from litellm.proxy.utils import get_available_models_for_user # Get available models for the user - all_models_str: Final = await get_available_models_for_user( - user_api_key_dict=user_api_key_dict, - llm_router=llm_router, - general_settings=general_settings, - user_model=user_model, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - team_id=None, - include_model_access_groups=False, - only_model_access_groups=False, - return_wildcard_routes=False, - user_api_key_cache=user_api_key_cache, + is_proxy_admin: Final = user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + all_models_str: Final = ( + llm_router.get_model_names() + if is_proxy_admin + else await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + general_settings=general_settings, + user_model=user_model, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + team_id=None, + include_model_access_groups=False, + only_model_access_groups=False, + return_wildcard_routes=False, + user_api_key_cache=user_api_key_cache, + ) ) model_groups: list[ModelGroupInfoProxy] = _get_model_group_info( llm_router=llm_router, all_models_str=all_models_str, model_group=model_group diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 4c141bcf698..a076b3593b1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -9,7 +9,7 @@ Pins (PR2): from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -128,7 +128,6 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path): assert "LLM Model List not loaded" in response.text - def test_get_proxy_model_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``GET /v1/model/info`` enriches each deployment through ``_get_proxy_model_info``; a registry entry declaring parallel function calling must land in ``model_info`` instead of null.""" @@ -161,9 +160,7 @@ def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch router.get_model_list = MagicMock(return_value=[deployment]) monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) - expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info( - [deployment] - ) + expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) allowed_model_names = proxy_server._get_v1_model_info_allowed_model_names( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", @@ -308,6 +305,61 @@ def test_model_group_info_invalid_method(client, auth_as, null_router): assert len(response.content) > 0 +@pytest.fixture +def model_group_info_router(monkeypatch): + from litellm.types.proxy.management_endpoints.model_management_endpoints import ModelGroupInfoProxy + + model_names = ["gpt-4", "claude-3"] + router = MagicMock() + router.get_model_names.return_value = model_names + router.get_model_access_groups.return_value = {} + router.get_model_list.return_value = [] + + def model_group_info(*, llm_router, all_models_str, model_group): + return [ModelGroupInfoProxy(model_group=name, providers=[]) for name in all_models_str] + + async def append_agents_to_model_group(*, model_groups, user_api_key_dict): + return model_groups + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": name} for name in model_names]) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", None) + monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info) + + from litellm.proxy.agent_endpoints import model_list_helpers + + monkeypatch.setattr( + model_list_helpers, + "append_agents_to_model_group", + AsyncMock(side_effect=append_agents_to_model_group), + ) + return router + + +def test_model_group_info_proxy_admin_ignores_key_model_restriction(client, auth_as, model_group_info_router): + from litellm.proxy._types import LitellmUserRoles + + with auth_as(LitellmUserRoles.PROXY_ADMIN, models=["no-default-models"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"] + + +def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router): + from litellm.proxy._types import LitellmUserRoles + + with auth_as(LitellmUserRoles.INTERNAL_USER, models=["gpt-4"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4"] + + # --------------------------------------------------------------------------- # GET /v2/model/info?exclude_auto_routers # --------------------------------------------------------------------------- @@ -399,14 +451,10 @@ def test_v2_model_info_exclude_auto_routers_shrinks_total_count(client, auth_as, assert len(payload["data"]) == payload["total_count"] -def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set( - client, auth_as, mixed_auto_router_router -): +def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(client, auth_as, mixed_auto_router_router): """Page size applies to the filtered list, so no page silently comes back short.""" with auth_as(): - response = client.get( - "/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1} - ) + response = client.get("/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1}) payload = response.json() assert payload["total_count"] == 2 assert payload["total_pages"] == 2 From e7454e52e6b15c9d7cce043abb84d43eecb1f0d7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:09:05 +0000 Subject: [PATCH 10/89] fix(gemini): map minimal thinking to low for Gemini 3.7 and 3.8 Flash Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 35 ++++----- ...test_vertex_and_google_ai_studio_gemini.py | 78 +++++++++++++++++++ 2 files changed, 95 insertions(+), 18 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d113b2b4f6b..d719d53e19f 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -110,6 +110,7 @@ else: SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") +GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING: Final = ("gemini-3.7-flash", "gemini-3.8-flash") def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: @@ -860,6 +861,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: raise _unsupported_reasoning_effort(reasoning_effort) + @staticmethod + def _supports_minimal_thinking_level(model: str) -> bool: + lowered: Final = model.lower() + is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered + return is_gemini3flash and not any(m in lowered for m in GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING) + @staticmethod def _map_reasoning_effort_to_thinking_level( reasoning_effort: str, @@ -874,13 +881,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Returns: GeminiThinkingConfig with thinkingLevel and includeThoughts """ - # Check if this is gemini-3-flash which supports MINIMAL thinking level - # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, - # gemini-3.5-flash, and any future 3.x-flash variants. is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower()) + supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model) is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": - if is_gemini3flash: + if supports_minimal: return {"thinkingLevel": "minimal", "includeThoughts": True} else: return {"thinkingLevel": "low", "includeThoughts": True} @@ -893,18 +898,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return {"thinkingLevel": "high", "includeThoughts": True} elif reasoning_effort == "high": return {"thinkingLevel": "high", "includeThoughts": True} - elif reasoning_effort == "disable": - # Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others - if is_gemini3flash: - return {"thinkingLevel": "minimal", "includeThoughts": False} - else: - return {"thinkingLevel": "low", "includeThoughts": False} - elif reasoning_effort == "none": - # For gemini-3-flash-preview, use "minimal" instead of "low" - if is_gemini3flash: - return {"thinkingLevel": "minimal", "includeThoughts": False} - else: - return {"thinkingLevel": "low", "includeThoughts": False} + elif reasoning_effort in ("disable", "none"): + return { + "thinkingLevel": "minimal" if supports_minimal else "low", + "includeThoughts": False, + } else: raise _unsupported_reasoning_effort(reasoning_effort) @@ -971,8 +969,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = True # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: - is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower() - params["thinkingLevel"] = "minimal" if is_gemini3flash else "low" + params["thinkingLevel"] = ( + "minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low" + ) else: # Thinking disabled params["includeThoughts"] = False diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 101f6e6fa5d..a1c31689d09 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2678,6 +2678,84 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): assert result["thinkingConfig"]["includeThoughts"] is False +@pytest.mark.parametrize( + "model", + [ + "gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", + "gemini-3.8-flash-preview", + ], +) +@pytest.mark.parametrize( + ("reasoning_effort", "include_thoughts"), + [("minimal", True), ("none", False), ("disable", False)], +) +def test_gemini_37_38_flash_floor_minimal_thinking_level( + model, reasoning_effort, include_thoughts +): + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + reasoning_effort, model + ) + + assert result["thinkingLevel"] == "low" + assert result["includeThoughts"] is include_thoughts + + +@pytest.mark.parametrize( + ("model", "reasoning_effort", "expected_level", "include_thoughts"), + [ + ("gemini-3-flash-preview", "minimal", "minimal", True), + ("gemini-3-flash-preview", "none", "minimal", False), + ("gemini-3-flash-preview", "disable", "minimal", False), + ("gemini-3.6-flash", "minimal", "minimal", True), + ("gemini-3.6-flash", "none", "minimal", False), + ("gemini-3.6-flash", "disable", "minimal", False), + ("gemini-3.5-flash", "minimal", "minimal", True), + ("gemini-3.5-flash", "none", "minimal", False), + ("gemini-3.5-flash", "disable", "minimal", False), + ("gemini-3.8-flash", "medium", "medium", True), + ], +) +def test_gemini_flash_minimal_thinking_support( + model, reasoning_effort, expected_level, include_thoughts +): + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + reasoning_effort, model + ) + + assert result["thinkingLevel"] == expected_level + assert result["includeThoughts"] is include_thoughts + + +def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): + monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) + thinking_param = {"type": "enabled", "budget_tokens": 1024} + + result_38 = VertexGeminiConfig._map_thinking_param( + thinking_param, model="gemini-3.8-flash" + ) + result_36 = VertexGeminiConfig._map_thinking_param( + thinking_param, model="gemini-3.6-flash" + ) + + assert result_38["thinkingLevel"] == "low" + assert result_36["thinkingLevel"] == "minimal" + + +def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): + result = VertexGeminiConfig().map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + assert result["thinkingConfig"] == { + "thinkingLevel": "low", + "includeThoughts": False, + } + + def test_reasoning_effort_dict_format_gemini_3(): """ Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. From 13553473aaf1afaac6157ff8d0c87a1984660c8e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:06:00 -0700 Subject: [PATCH 11/89] fix(mcp): reject missing upstream authentication credentials --- litellm/experimental_mcp_client/client.py | 8 +- .../mcp_server/mcp_server_manager.py | 87 ++++---- .../mcp_server/openapi_to_mcp_generator.py | 24 +-- .../outbound_credentials/adapter.py | 24 +-- .../proxy/_experimental/mcp_server/server.py | 1 + .../_experimental/mcp_server/upstream.py | 81 ++++++++ .../proxy/_experimental/mcp_server/utils.py | 16 ++ .../test_mcp_client.py | 15 ++ .../outbound_credentials/test_adapter.py | 17 +- .../mcp_server/test_mcp_hook_extra_headers.py | 1 + .../mcp_server/test_mcp_server_manager.py | 190 +++++++++++++++++- .../test_openapi_to_mcp_generator.py | 18 ++ 12 files changed, 396 insertions(+), 86 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/upstream.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ee01a53ecb3..56ee5f30d02 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -346,13 +346,17 @@ class MCPClient: self.update_auth_value(auth_value) async def discovery_auth_fingerprint(self) -> str: + return self._hash_discovery_auth(await self.prepare_request_auth()) + + async def prepare_request_auth(self) -> httpx.Request: + """Preview the authenticated request without sending it, closing the auth flow afterwards.""" request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers()) if self._resolved_auth is None: - return self._hash_discovery_auth(request) + return request flow: Final = self._resolved_auth.async_auth_flow(request) try: authenticated: Final = await flow.__anext__() - return self._hash_discovery_auth(authenticated) + return authenticated finally: await flow.aclose() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fb0c623473a..0254f79cbcc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -132,6 +132,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) +from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client, validate_openapi_credentials from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -4229,16 +4230,19 @@ class MCPServerManager: ) record_auth_resolution(server.server_id, AuthResolution.not_applicable) - return MCPClient( - server_url="", # Not used for stdio - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - stdio_config=stdio_config, - extra_headers=extra_headers, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + stdio_config=stdio_config, + extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) else: # For HTTP/SSE transports @@ -4259,15 +4263,20 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=extra_headers, ) - return MCPClient( - server_url=server_url, - transport_type=transport, - auth_type=resolved_server.auth_type, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - extra_headers=extra_headers, - resolved_auth=resolved_auth, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=resolved_server.auth_type, + timeout=( + resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT + ), + extra_headers=extra_headers, + resolved_auth=resolved_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) # Create SigV4 auth if configured @@ -4297,17 +4306,20 @@ class MCPServerManager: else AuthResolution.no_auth ) record_auth_resolution(server.server_id, legacy_source) - return MCPClient( - server_url=server_url, - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - auth_header_name=auth_header_name, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - extra_headers=extra_headers, - aws_auth=aws_auth, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + auth_header_name=auth_header_name, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + extra_headers=extra_headers, + aws_auth=aws_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) async def _get_tools_from_server( @@ -6188,6 +6200,7 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, forwarded_headers: dict[str, str] | None, + caller_authorization: str | None = None, ) -> tuple[dict[str, str] | None, dict[str, str] | None]: """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. @@ -6211,9 +6224,12 @@ class MCPServerManager: """ spec: Final = to_server_spec(mcp_server) if spec is None: - if oauth2_headers: - return None, forwarded_headers - stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) + stored_headers = ( + None + if oauth2_headers + else await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) + ) + validate_openapi_credentials(mcp_server, stored_headers, forwarded_headers, caller_authorization) return stored_headers, forwarded_headers subject_token: str | None = None @@ -6232,7 +6248,9 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=forwarded_headers, ) - return await _materialize_auth_headers(resolved_auth), forwarded_headers + resolved_headers: Final = await _materialize_auth_headers(resolved_auth) + validate_openapi_credentials(mcp_server, resolved_headers, forwarded_headers, caller_authorization) + return resolved_headers, forwarded_headers async def _gather_openapi_tool_tasks( self, @@ -6358,6 +6376,7 @@ class MCPServerManager: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, + caller_authorization=auth_header_value, ) async def _call_openapi_via_handler(): diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index d115eb8b3c1..66712e97a34 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -20,6 +20,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, MCPUpstreamAuthError, ) +from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to @@ -415,26 +416,9 @@ def _merge_openapi_tool_request_headers( Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ - request_extra: Final = _request_extra_headers.get() or {} - static: Final = static_headers or {} - - static_lower_names: Final = {k.lower() for k in static} - effective_headers: dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} - effective_headers.update(static) - - override_auth: Final = _request_auth_header.get() - if override_auth: - for existing in [k for k in effective_headers if k.lower() == "authorization"]: - del effective_headers[existing] - effective_headers["Authorization"] = override_auth - - resolved_auth_headers: Final = _request_resolved_auth_headers.get() or {} - for name, value in resolved_auth_headers.items(): - for existing in [k for k in effective_headers if k.lower() == name.lower()]: - del effective_headers[existing] - effective_headers[name] = value - - return effective_headers + return merge_openapi_headers( + static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get() + ) def _raise_for_upstream_failure( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 77979a15199..d25946d81d0 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -79,7 +79,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None: BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers - to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later). + to v1 for its static schemes. Declared OBO always stays with the exchange arm. Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is @@ -90,8 +90,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None: modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough oauth2 and SigV4 return None and stay on v1. """ - if server.is_byok: - return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) + if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange: + return None # per-user BYOK source not migrated yet -> defer to v1 resource: Final = server.url or server.server_id auth_type: Final = server.auth_type match auth_type: @@ -165,21 +165,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: ) -def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: - """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured. - - An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the - ``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at - the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the - gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is - nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect - (``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value - normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is - forwarded only when the operator set it; a missing one is omitted, not derived. - """ +def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec: + """Keep declared OBO owned by the resolver, including incomplete client configuration.""" endpoint: Final = server.token_exchange_endpoint or server.effective_token_url - if not server.client_id or not server.client_secret: - return None profile: Final[Literal["rfc8693", "entra_obo"]] = ( "entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693" ) @@ -193,7 +181,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: token_exchange_endpoint=endpoint, audience=server.audience, client_id=server.client_id, - client_secret=SecretStr(server.client_secret), + client_secret=SecretStr(server.client_secret) if server.client_secret else None, token_endpoint_auth_method=server.token_endpoint_auth_method, scopes=tuple(server.scopes or ()), ), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7feb1fd468d..a3aaada41f7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3141,6 +3141,7 @@ if MCP_AVAILABLE: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, + caller_authorization=auth_header_value, ) _auth_token: Final = _request_auth_header.set(auth_header_value) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py new file mode 100644 index 00000000000..8d7a7dba252 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from typing import Final + +from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError +from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_STATIC_MODES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) +) + + +def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: + if not value: + return False + if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + return True + if value.lower() in ("bearer", "basic", "token", "apikey"): + return False + if auth_type == MCPAuth.basic: + parts: Final = value.split(None, 1) + if len(parts) != 2 or parts[0].lower() != "basic": + return False + try: + return bool(base64.b64decode(parts[1], validate=True).strip()) + except ValueError: + return False + return True + + +def validate_static_credential( + server: MCPServer, headers: Mapping[str, str], *, header_slot: str | None = None, openapi: bool = False +) -> Result[None, CredError]: + if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: + return Ok(None) + default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" + slots: Final = frozenset( + name.lower() + for name in ( + header_slot or server.upstream_token_header or default_slot, + "Authorization" if openapi else default_slot, + ) + ) + values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if values and all(_usable_credential_value(server.auth_type, name, value) for name, value in values): + return Ok(None) + return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) + + +async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: + if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: + return client + request: Final = await client.prepare_request_auth() + match validate_static_credential(server, request.headers): + case Error(error): + raise_public(error) + case Ok(): + return client + + +def validate_openapi_credentials( + server: MCPServer, + resolved_headers: Mapping[str, str] | None, + forwarded_headers: Mapping[str, str] | None, + caller_authorization: str | None, +) -> None: + headers: Final = merge_openapi_headers( + server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers + ) + match validate_static_credential(server, headers, openapi=True): + case Error(error): + raise_public(error) + case Ok(): + return diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index fb3eb06fd15..bea74d36b34 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -756,6 +756,22 @@ def build_env_var_setup_url(server_id: str) -> str: return f"{base}{path}" if base else path +def merge_openapi_headers( + static_headers: Mapping[str, str], + extra_headers: Mapping[str, str] | None, + caller_authorization: str | None, + resolved_headers: Mapping[str, str] | None, +) -> dict[str, str]: + sources: Final = ( + extra_headers or {}, + static_headers, + {"Authorization": caller_authorization} if caller_authorization else {}, + resolved_headers or {}, + ) + entries: Final = {name.lower(): (name, value) for source in sources for name, value in source.items()} + return dict(entries.values()) + + def merge_mcp_headers( *, extra_headers: Mapping[str, str] | None = None, 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 f72316f5d5e..d9ffb0d64fe 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1934,3 +1934,18 @@ async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: assert original != replaced assert len(original) == 64 assert "private-original-credential" not in original + + +@pytest.mark.asyncio +async def test_request_auth_preview_uses_the_same_effective_headers_as_egress() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + + client: Final = MCPClient( + server_url="https://upstream.example/mcp", auth_type=MCPAuth.bearer_token, + resolved_auth=StaticHeaderAuth("Bearer resolved"), extra_headers={"X-Trace": "trace"}, + ) + request: Final = await client.prepare_request_auth() + assert request.method == "POST" + assert str(request.url) == "https://upstream.example/mcp" + assert request.headers["Authorization"] == "Bearer resolved" + assert request.headers["X-Trace"] == "trace" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 1b003e11993..2885fdaef95 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -155,12 +155,6 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured _server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1 - _server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1 - _server( - auth_type=MCPAuth.oauth2_token_exchange, - token_exchange_endpoint="https://idp/token", - client_id="cid", - ), # missing client_secret -> incomplete -> v1 _server(auth_type=MCPAuth.aws_sigv4), _server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]), ], @@ -802,3 +796,14 @@ def test_a_blank_header_name_means_unset_rather_than_an_error(blank): spec = to_server_spec(server) assert spec is not None assert spec.config.header_name == "Authorization" + + +@pytest.mark.parametrize("client_secret", [None, ""]) +@pytest.mark.parametrize("is_byok", [False, True]) +def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_byok: bool) -> None: + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2_token_exchange, client_id="client", + client_secret=client_secret, is_byok=is_byok)) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.client_id == "client" + assert spec.config.client_secret is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 28faf375ab8..5e3a26fb4ac 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1391,6 +1391,7 @@ class TestOpenApiResolvedUpstreamAuth: mcp_auth_header="user-byok-key", user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), forwarded_headers=None, + caller_authorization="ApiKey user-byok-key", ) assert resolved is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d56f08c4e79..ea05035e16f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -9401,12 +9401,13 @@ class TestCreateMcpClientV2Graft: assert "misconfigured" in str(exc_info.value.detail) assert "token_url" in str(exc_info.value.detail) - async def test_static_token_missing_defers_to_v1(self): - client = await MCPServerManager()._create_mcp_client( - self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) - ) - - assert client._resolved_auth is None + async def test_static_token_missing_rejects_before_connecting(self): + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) + ) + assert exc.value.status_code == 500 + assert "credential" in str(exc.value.detail) async def test_stdio_migrated_auth_type_still_defers_to_v1(self): client = await MCPServerManager()._create_mcp_client( @@ -13467,3 +13468,180 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( result: Final = await cache.get(("server", None), fetch) assert result[0].description == description assert fetch.await_count == 2 + + +class TestProtectedCredentialPreparation: + @pytest.mark.asyncio + @pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse]) + @pytest.mark.parametrize("client_secret", [None, ""]) + @pytest.mark.parametrize("subject", [None, "caller-subject"]) + async def test_incomplete_obo_rejects_caller_and_static_fallback( + self, transport: MCPTransport, client_secret: str | None, subject: str | None + ) -> None: + server = MCPServer( + server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp", + transport=transport, auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway", client_secret=client_secret, + token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client( + server, mcp_auth_header="Bearer override", subject_token=subject, + ) + assert exc.value.status_code == (401 if subject is None else 500) + assert "static-fallback" not in str(exc.value.detail) + assert "override" not in str(exc.value.detail) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.api_key, MCPAuth.bearer_token]) + @pytest.mark.parametrize("credential", [None, "", " ", {"X-Trace": "trace"}]) + async def test_static_auth_without_usable_credential_rejects( + self, auth_type: MCPAuthType, credential: str | dict[str, str] | None + ) -> None: + server = MCPServer( + server_id="empty-static", name="empty-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential) + assert exc.value.status_code == 500 + assert "credential" in str(exc.value.detail).lower() + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,headers", [ + (MCPAuth.api_key, {"X-API-Key": "key"}), + (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), + ]) + async def test_static_auth_accepts_actual_forwarded_credential( + self, auth_type: MCPAuthType, headers: dict[str, str] + ) -> None: + server = MCPServer( + server_id="header-static", name="header-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + ) + client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers) + assert client._get_auth_headers() == headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token]) + async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: + server = MCPServer( + server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + token_exchange_endpoint="https://idp.example/token", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, + user_api_key_auth=None, forwarded_headers=None, + ) + assert exc.value.status_code in (401, 500) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,slot", [(MCPAuth.api_key, "X-API-Key"), (MCPAuth.authorization, "Authorization")]) + async def test_raw_static_value_named_token_is_a_usable_credential(self, auth_type: MCPAuthType, slot: str) -> None: + server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token="token") + client = await MCPServerManager()._create_mcp_client(server) + assert client._resolved_auth is not None + request = httpx.Request("GET", server.url) + flow = client._resolved_auth.auth_flow(request) + try: + assert next(flow).headers[slot] == "token" + finally: + flow.close() + + @pytest.mark.asyncio + async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: + server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True, + token_exchange_endpoint="https://idp.example/token") + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override") + assert exc.value.status_code == 401 + + @pytest.mark.asyncio + @pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")]) + async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None: + server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured) + client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override) + assert client._get_auth_headers()["Authorization"] == override + + @pytest.mark.asyncio + @pytest.mark.parametrize("token", [None, "shared"]) + async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None: + server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "}) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + async def test_custom_slot_uses_its_actual_credential(self) -> None: + server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", authentication_token="key") + client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"}) + assert client._credential_slot == "X-Custom" + assert await client.discovery_auth_fingerprint() + + @pytest.mark.asyncio + @pytest.mark.parametrize("static,forwarded,caller", [ + ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), + ({}, {"X-API-Key": "forwarded"}, None), + ({}, None, "ApiKey caller"), + ]) + async def test_openapi_static_credentials_remain_supported( + self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + ) -> None: + server = MCPServer(server_id="openapi-static", name="openapi-static", url="https://upstream.example", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static) + resolved, retained = await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, + user_api_key_auth=None, forwarded_headers=forwarded, caller_authorization=caller, + ) + assert resolved is None + assert retained == forwarded + + @pytest.mark.asyncio + async def test_static_resolution_cancellation_closes_flow(self) -> None: + from collections.abc import AsyncGenerator + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client + + class CancelledAuth(httpx.Auth): + closed = False + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + try: + raise asyncio.CancelledError() + yield request + finally: + self.closed = True + + auth = CancelledAuth() + server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key) + client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth) + with pytest.raises(asyncio.CancelledError): + await prepare_mcp_client(server, client) + assert auth.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization]) + async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None: + server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ") + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc"]) + async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: + server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) + assert exc.value.status_code == 500 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 5fa202224e3..66c5627bc94 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1458,3 +1458,21 @@ class TestBoundedOpenAPISpecLoading: else: assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} assert destination.call_count == 1 + + +def test_openapi_generator_import_does_not_require_mcp_sdk() -> None: + import subprocess + import sys + + script = """ +import builtins +original_import = builtins.__import__ +def without_mcp(name, *args, **kwargs): + if name == 'mcp' or name.startswith('mcp.'): + raise ModuleNotFoundError('MCP SDK unavailable') + return original_import(name, *args, **kwargs) +builtins.__import__ = without_mcp +import litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) + assert result.returncode == 0, result.stderr From 8a3add3c6a7ad900e19d2ed7760627612aa3d17d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:27:25 -0700 Subject: [PATCH 12/89] fix(mcp): reject scheme-only Basic credentials --- litellm/proxy/_experimental/mcp_server/upstream.py | 3 ++- .../mcp_server/test_mcp_server_manager.py | 13 ++++++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 8d7a7dba252..89fd1ad5066 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -29,7 +29,8 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b if len(parts) != 2 or parts[0].lower() != "basic": return False try: - return bool(base64.b64decode(parts[1], validate=True).strip()) + decoded: Final = base64.b64decode(parts[1], validate=True).strip() + return bool(decoded) and decoded.lower() != b"basic" except ValueError: return False return True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ea05035e16f..e2275d05fcf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13638,10 +13638,21 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc"]) + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", transport=MCPTransport.http, auth_type=MCPAuth.basic) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["Basic", "Basic ", "basic"]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_basic_scheme_alone_is_not_a_credential(self, value: str, source: str) -> None: + server = MCPServer(server_id="basic-scheme", name="basic-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 From 84e14789d1c7861fbfe76534c6eb091bcaf00f7b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:18:48 -0700 Subject: [PATCH 13/89] fix(mcp): preserve usable alternate header credentials --- .../_experimental/mcp_server/upstream.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 89fd1ad5066..49aca5cf31b 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -50,7 +50,7 @@ def validate_static_credential( ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) - if values and all(_usable_credential_value(server.auth_type, name, value) for name, value in values): + if any(_usable_credential_value(server.auth_type, name, value) for name, value in values): return Ok(None) return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e2275d05fcf..c8dba1c1554 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13591,6 +13591,7 @@ class TestProtectedCredentialPreparation: ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), ({}, {"X-API-Key": "forwarded"}, None), ({}, None, "ApiKey caller"), + ({"X-API-Key": "static"}, {"Authorization": ""}, None), ]) async def test_openapi_static_credentials_remain_supported( self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None @@ -13656,3 +13657,39 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value,default_slot", [ + (MCPAuth.api_key, "fixture-key", "X-API-Key"), + (MCPAuth.bearer_token, "fixture-key", "Authorization"), + (MCPAuth.basic, "user:pass", "Authorization"), + (MCPAuth.token, "fixture-key", "Authorization"), + (MCPAuth.authorization, "fixture-key", "Authorization"), + ]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_usable_credential_survives_an_empty_alternate_header( + self, auth_type: MCPAuthType, value: str, default_slot: str, source: str + ) -> None: + server: Final = MCPServer( + server_id="alternate", name="alternate", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, upstream_token_header="X-Custom", + authentication_token=value if source == "configured" else None, + ) + empty_slot: Final = default_slot if source == "configured" else "X-Custom" + selected_slot: Final = "X-Custom" if source == "configured" else default_slot + client: Final = await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=value if source == "caller" else None, extra_headers={empty_slot: ""}, + ) + request: Final = await client.prepare_request_auth() + assert request.headers[selected_slot] + assert request.headers[empty_slot] == "" + + @pytest.mark.asyncio + async def test_empty_custom_and_default_headers_do_not_satisfy_auth(self) -> None: + server: Final = MCPServer( + server_id="both-empty", name="both-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) + assert exc.value.status_code == 500 From 1a7ca04cc5b908c37b2b78efd1cf683bfdd2163b Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 02:29:26 +0000 Subject: [PATCH 14/89] fix(proxy): carry litellm_call_id through endpoint specific error logs and failure responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 10 ++- litellm/proxy/batches_endpoints/endpoints.py | 27 ++++--- litellm/proxy/common_request_processing.py | 41 +++++++--- .../common_utils/openai_error_payload.py | 6 ++ litellm/proxy/image_endpoints/endpoints.py | 24 +++--- .../pass_through_endpoints.py | 12 ++- litellm/proxy/proxy_server.py | 46 ++++++----- litellm/proxy/rerank_endpoints/endpoints.py | 20 +++-- litellm/proxy/utils.py | 9 ++- .../anthropic_endpoints/test_endpoints.py | 81 +++++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 23 ++++++ .../common_utils/test_openai_error_payload.py | 6 ++ .../proxy/image_endpoints/test_endpoints.py | 68 ++++++++++++++++ .../test_pass_through_endpoints.py | 39 +++++++++ .../proxy/rerank_endpoints/test_endpoints.py | 51 ++++++++++-- .../proxy/test_common_request_processing.py | 8 +- tests/test_litellm/proxy/test_proxy_server.py | 43 ++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 31 +++++++ .../proxy/utils/helpers/test_error_helpers.py | 15 ++++ 19 files changed, 485 insertions(+), 75 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index d4cb3b84ee4..f673a82654d 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse import litellm -from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.anthropic.experimental_pass_through.context_management import ( @@ -22,7 +21,9 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, create_response, + log_llm_api_exception, proxy_exception_from_http_exception, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_error_payload import ( @@ -218,7 +219,7 @@ async def anthropic_response( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e) + log_llm_api_exception(e, base_llm_response_processor.litellm_call_id) if isinstance(e, ProxyException): return _anthropic_error_json_response(e, request) @@ -231,7 +232,7 @@ async def anthropic_response( # Get headers headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=data.get("litellm_call_id", ""), + call_id=base_llm_response_processor.litellm_call_id, model_id=model_id, version=version, response_cost=0, @@ -288,6 +289,7 @@ async def count_tokens( """ from litellm.proxy.proxy_server import token_counter as internal_token_counter + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: request_data: Final = await _read_request_body(request=request) data: Final[dict] = {**request_data} @@ -339,7 +341,7 @@ async def count_tokens( detail=detail, ) except Exception as e: - verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e) + log_llm_api_exception(e, litellm_call_id) raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..e3767d06e7d 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -17,7 +17,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + request_litellm_call_id, +) from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_endpoint_utils import ( @@ -383,8 +387,9 @@ async def create_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.get( @@ -674,8 +679,9 @@ async def retrieve_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.get( @@ -725,6 +731,7 @@ async def list_batches( ) verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit) + data: dict = {} try: if llm_router is None: raise HTTPException( @@ -856,8 +863,9 @@ async def list_batches( original_exception=e, request_data={"after": after, "limit": limit}, ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.post( @@ -1079,8 +1087,9 @@ async def cancel_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) ###################################################################### diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 46b222a4fc9..ba8f48b88be 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -7,7 +7,18 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NamedTuple, + Protocol, + TypeAlias, + TypeVar, + overload, + runtime_checkable, +) import anyio import httpx @@ -1452,7 +1463,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool: _CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request" -def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None: +@runtime_checkable +class _CarriesLitellmCallId(Protocol): + litellm_call_id: str | None + + +def request_litellm_call_id(data: Mapping[str, object]) -> str | None: + logging_obj: Final = data.get("litellm_logging_obj") + logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None + call_id: Final = logged_id or data.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None: if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL: verbose_proxy_logger.info( "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, " @@ -1532,6 +1555,10 @@ class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data + @property + def litellm_call_id(self) -> str | None: + return request_litellm_call_id(self.data) + @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, @@ -3429,11 +3456,7 @@ class ProxyBaseLLMRequestProcessing: version: str | None = None, ): """Raises ProxyException (OpenAI API compatible) if an exception is raised""" - logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None) - _log_llm_api_exception( - e, - (logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"), - ) + log_llm_api_exception(e, self.litellm_call_id) # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -3463,9 +3486,7 @@ class ProxyBaseLLMRequestProcessing: custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=( - _litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id") - ), + call_id=self.litellm_call_id, model_id=model_id, version=version, response_cost=0, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index fe23ab2c4b6..d4312d93559 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -52,3 +52,9 @@ def openai_error_param(exc: object) -> str | None: serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None + + +def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers + if litellm_call_id is None: + return None + return {"x-litellm-call-id": litellm_call_id} # mutable-ok: ProxyException mutates its headers dict diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 3f044855ce8..30406bbcaae 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,6 +1,5 @@ import asyncio import io -import traceback from collections.abc import Sequence from typing import Final, get_type_hints @@ -9,19 +8,23 @@ from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, from fastapi.responses import ORJSONResponse import litellm -from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import ( coerce_numeric_form_fields, numeric_form_fields, ) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -92,6 +95,7 @@ async def image_generation( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -106,6 +110,7 @@ async def image_generation( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if isinstance(model, str): reject_url_valued_destination("model", model) @@ -153,9 +158,7 @@ async def image_generation( response = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.) response = await proxy_logging_obj.post_call_success_hook( @@ -168,7 +171,7 @@ async def image_generation( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -179,7 +182,7 @@ async def image_generation( version=version, response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, ) @@ -200,13 +203,13 @@ async def image_generation( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.image_generation(): Exception occured - %s", e) - verbose_proxy_logger.debug(traceback.format_exc()) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: @@ -215,6 +218,7 @@ async def image_generation( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..e93b1232836 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -72,7 +72,9 @@ from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_end from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, + log_llm_api_exception, open_sse_before_first_byte, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -80,6 +82,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -197,6 +200,7 @@ async def chat_completion_pass_through_endpoint( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: body: Final = await request.body() body_str: Final = body.decode() @@ -224,6 +228,7 @@ async def chat_completion_pass_through_endpoint( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id # override with user settings, these are params passed via cli if user_temperature: @@ -290,9 +295,7 @@ async def chat_completion_pass_through_endpoint( response_cost: Final = hidden_params.get("response_cost", None) or "" ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) verbose_proxy_logger.debug("final response: %s", response) @@ -313,12 +316,13 @@ async def chat_completion_pass_through_endpoint( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..de42dcec8d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -348,7 +348,10 @@ from litellm.proxy.common_request_processing import ( _is_azure_model_router_request, _should_return_raw_model_name, create_response, + log_llm_api_exception, open_sse_before_first_byte, + request_litellm_call_id, + resolve_litellm_call_id, ttft_keepalive_interval, ) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -387,6 +390,7 @@ from litellm.proxy.common_utils.model_listing_utils import ( from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) +from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers from litellm.proxy.common_utils.periodic_reload_schedule import ( MODEL_COST_MAP_RELOAD_PARAM_NAME, clear_reload_interval, @@ -11291,12 +11295,14 @@ async def completion( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) @@ -11454,6 +11460,7 @@ async def moderations( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -11468,6 +11475,7 @@ async def moderations( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id data["model"] = ( general_settings.get("moderation_model", None) # server default @@ -11494,9 +11502,7 @@ async def moderations( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11522,7 +11528,7 @@ async def moderations( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.moderations(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): raise if isinstance(e, HTTPException): @@ -11530,6 +11536,7 @@ async def moderations( message=getattr(e, "message", str(e)), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: @@ -11538,6 +11545,7 @@ async def moderations( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", 500), ) @@ -11576,6 +11584,7 @@ async def audio_speech( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -11590,6 +11599,7 @@ async def audio_speech( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11612,9 +11622,7 @@ async def audio_speech( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11622,7 +11630,7 @@ async def audio_speech( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, @@ -11633,7 +11641,7 @@ async def audio_speech( response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), fastest_response_batch_completion=None, - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, ) @@ -11669,14 +11677,14 @@ async def audio_speech( original_exception=e, request_data=data, ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e) - verbose_proxy_logger.debug(traceback.format_exc()) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, (ProxyException, HTTPException)): raise e raise ProxyException( message=getattr(e, "message", f"{e}"), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) @@ -11705,6 +11713,7 @@ async def audio_transcriptions( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly form_data: Final = await get_form_data(request) @@ -11719,6 +11728,7 @@ async def audio_transcriptions( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11775,9 +11785,7 @@ async def audio_transcriptions( file_object.close() # close the file read in by io library ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11785,7 +11793,7 @@ async def audio_transcriptions( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" additional_headers: Final[dict] = hidden_params.get("additional_headers", {}) or {} fastapi_response.headers.update( @@ -11797,7 +11805,7 @@ async def audio_transcriptions( version=version, response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, **additional_headers, @@ -11819,12 +11827,13 @@ async def audio_transcriptions( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.audio_transcription(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: @@ -11833,6 +11842,7 @@ async def audio_transcriptions( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 16cd7368e4a..4f5eb411e44 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -7,12 +7,16 @@ import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -55,6 +59,7 @@ async def rerank( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: body: Final = await request.body() data = orjson.loads(body) @@ -68,6 +73,7 @@ async def rerank( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") @@ -82,9 +88,7 @@ async def rerank( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -95,7 +99,7 @@ async def rerank( fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None), + call_id=hidden_params.get("litellm_call_id", None) or litellm_call_id, model_id=model_id, cache_key=cache_key, api_base=api_base, @@ -113,12 +117,13 @@ async def rerank( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.rerank(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: @@ -127,5 +132,6 @@ async def rerank( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 479bd0a55af..53782227998 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,7 +38,7 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) -from litellm.proxy.common_utils.openai_error_payload import openai_error_param +from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers, openai_error_param from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -3031,7 +3031,7 @@ class ProxyLogging: if litellm_logging_obj is None: from litellm._uuid import uuid - request_data["litellm_call_id"] = str(uuid.uuid4()) + request_data.setdefault("litellm_call_id", str(uuid.uuid4())) user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( user_api_key_dict=user_api_key_dict ) @@ -7638,7 +7638,7 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")) -def handle_exception_on_proxy(e: Exception) -> ProxyException: +def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible """ @@ -7650,11 +7650,13 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: _recreate_writer_on_read_only_transaction(prisma_client) + headers: Final = litellm_call_id_headers(litellm_call_id) if isinstance(e, HTTPException): return ProxyException( message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, param=openai_error_param(e), + headers=headers, code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): @@ -7664,6 +7666,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: message=str(e), type=ProxyErrorTypes.internal_server_error, param=openai_error_param(e), + headers=headers, code=_status_code, ) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index f809fadc879..e4b15cfdcd3 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -3,12 +3,14 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec """ import json +import logging import unittest from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -285,6 +287,85 @@ class TestFailureHookRequestData: assert hook_request_data["litellm_logging_obj"] == "logging-obj-sentinel" +class TestErrorLogCarriesCallId: + """LIT-7836: the /v1/messages and /v1/messages/count_tokens error lines must carry + the request's litellm_call_id, rendered in the message and as a structured field.""" + + @pytest.fixture(autouse=True) + def propagating_proxy_logger(self): + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + @staticmethod + def _error_record(caplog: pytest.LogCaptureFixture) -> logging.LogRecord: + return next(r for r in caplog.records if "Exception occured" in r.getMessage()) + + @pytest.mark.asyncio + async def test_messages_failure_log_carries_call_id(self, caplog: pytest.LogCaptureFixture): + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + call_id = "messages-call-7836" + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_call_id": call_id} + raise RuntimeError("provider timeout") + + request = MagicMock() + request.headers = {} + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + mock_logging.post_call_failure_hook = AsyncMock() + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.status_code == 500 + record = self._error_record(caplog) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio + async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): + from fastapi import HTTPException + + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + call_id = "count-tokens-call-7836" + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + + with ( + patch.object( # test-quality-ok: endpoint reads the body via a module function; no injection seam + ep, + "_read_request_body", + new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), + ), + patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + pytest.raises(HTTPException) as raised, + ): + await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.status_code == 500 + record = self._error_record(caplog) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..cf805b384d2 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -31,6 +31,7 @@ cannot drift without a test failure. import base64 import json +import logging from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict, Optional @@ -1088,6 +1089,28 @@ async def test_create__exception_calls_failure_hook(harness, openai_env_creds): assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" +async def test_create__exception_carries_the_litellm_call_id(harness, openai_env_creds, caplog): + call_id = "lit7836-batch-call-id" + set_body( + harness, + { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "litellm_call_id": call_id, + }, + ) + harness.litellm_acreate.side_effect = ValueError("provider boom") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await call_create(harness) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + # =========================================================================== # # # # GET /v1/batches/{batch_id} - retrieve_batch routing-contract tests # diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 90850840ab4..df775916046 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -6,6 +6,7 @@ from fastapi import HTTPException from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -158,3 +159,8 @@ def test_a_stringified_none_type_or_param_is_treated_as_absent(): assert carried.type == "None" assert openai_error_type(carried, 400) == "invalid_request_error" assert openai_error_param(carried) is None + + +def test_a_failed_request_answers_with_the_call_id_it_was_logged_under(): + assert litellm_call_id_headers("call-7836") == {"x-litellm-call-id": "call-7836"} + assert litellm_call_id_headers(None) is None diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index d8b3eef98bd..d03832bf6d0 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -1,5 +1,7 @@ import asyncio import copy +import logging +from collections.abc import Iterator from types import SimpleNamespace from typing import Any, Dict @@ -10,6 +12,7 @@ from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.image_endpoints import endpoints @@ -211,3 +214,68 @@ async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(mon await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404") + + +@pytest.fixture +def propagating_proxy_logger() -> Iterator[None]: + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + +@pytest.mark.asyncio +async def test_failure_log_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None +) -> None: + """LIT-7836: the /v1/images/generations error line must carry the litellm_call_id + the client sent, both rendered in the message and as a structured record field.""" + call_id = "images-call-7836" + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]: + return data + + async def fake_post_call_failure_hook(**_: object) -> None: + return None + + async def failing_route_request(**_: object) -> None: + raise HTTPException(status_code=401, detail={"error": "invalid api key"}) + + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", + SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version") + monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request) + + body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"}) + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/images/generations", + "headers": [(b"x-litellm-call-id", call_id.encode())], + }, + receive, + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0fc961cf8c9..f5348c8adc1 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -6021,3 +6021,42 @@ async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_err ) assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "400") + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_failure_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + call_id = "lit7836-pass-through-call-id" + proxy_logging = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + request = MagicMock(spec=Request) + request.headers = Headers({"x-litellm-call-id": call_id}) + request.body = AsyncMock( + return_value=json.dumps({"model": "unknown-model", "messages": [{"role": "user", "content": "hi"}]}).encode() + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py index ea858e04e0f..52d12dd1813 100644 --- a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py @@ -3,6 +3,8 @@ Tests for rerank_endpoints/endpoints.py response headers. """ import json +import logging +from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,6 +12,7 @@ from fastapi import HTTPException, Request, Response import litellm.proxy.common_request_processing as common_request_processing_mod import litellm.proxy.proxy_server as proxy_server_mod +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.rerank_endpoints.endpoints import rerank from litellm.types.utils import RerankResponse @@ -28,7 +31,7 @@ HIDDEN_PARAMS = { } -def _build_request() -> Request: +def _build_request(headers: tuple[tuple[bytes, bytes], ...] = ()) -> Request: body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode() async def receive(): @@ -39,7 +42,7 @@ def _build_request() -> Request: "type": "http", "method": "POST", "path": "/rerank", - "headers": [(b"content-type", b"application/json")], + "headers": [(b"content-type", b"application/json"), *headers], "query_string": b"", }, receive=receive, @@ -56,7 +59,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: proxy_logging_obj.update_request_status = AsyncMock() async def fake_add_litellm_data_to_request(**kwargs): - return {**kwargs["data"], "litellm_call_id": "call-123"} + return dict(kwargs["data"]) async def fake_route_request(**kwargs): async def _call(): @@ -72,7 +75,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler ): await rerank( - request=_build_request(), + request=_build_request(headers=((b"x-litellm-call-id", b"call-123"),)), fastapi_response=fastapi_response, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) @@ -121,7 +124,11 @@ async def test_rerank_omits_detailed_timing_headers_when_disabled(): async def _rerank_failure( - failure: Exception, *, raised_before_routing: bool, monkeypatch: pytest.MonkeyPatch + failure: Exception, + *, + raised_before_routing: bool, + monkeypatch: pytest.MonkeyPatch, + headers: tuple[tuple[bytes, bytes], ...] = (), ) -> ProxyException: proxy_logging_obj = MagicMock() proxy_logging_obj.pre_call_hook = AsyncMock( @@ -143,13 +150,45 @@ async def _rerank_failure( with pytest.raises(ProxyException) as raised: await rerank( - request=_build_request(), + request=_build_request(headers), fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) return raised.value +@pytest.fixture +def propagating_proxy_logger() -> Iterator[None]: + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + +@pytest.mark.asyncio +async def test_failure_log_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None +) -> None: + """LIT-7836: the /rerank error line must carry the same litellm_call_id the client + sent, both in the rendered message and as a structured log record field.""" + call_id = "rerank-call-7836" + failure = HTTPException(status_code=401, detail={"error": "invalid api key"}) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + raised = await _rerank_failure( + failure, + raised_before_routing=False, + monkeypatch=monkeypatch, + headers=((b"x-litellm-call-id", call_id.encode()),), + ) + + assert raised.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch): """A bare HTTPException carries no type or param, so the tail used to ship the diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 099204cd6c6..a621a97d448 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8212,7 +8212,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ """Regression for LIT-6043: expected 4xx errors log without formatting a traceback; unexpected errors keep logger.exception behavior.""" from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_request_processing import _log_llm_api_exception + from litellm.proxy.common_request_processing import log_llm_api_exception verbose_proxy_logger.propagate = True try: @@ -8220,7 +8220,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ try: raise exc except Exception as raised: - _log_llm_api_exception(raised, "call-id-for-traceback-test") + log_llm_api_exception(raised, "call-id-for-traceback-test") finally: verbose_proxy_logger.propagate = False @@ -8778,14 +8778,14 @@ class TestErrorLogCarriesCallId: from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ( _CLIENT_DISCONNECT_DETAIL, - _log_llm_api_exception, + log_llm_api_exception, ) call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.propagate = True try: with caplog.at_level("INFO", logger="LiteLLM Proxy"): - _log_llm_api_exception( + log_llm_api_exception( HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL), call_id, ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6f55449abab..85c1a23c473 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2,6 +2,7 @@ import asyncio import contextlib import importlib import json +import logging import os import re import socket @@ -12927,6 +12928,48 @@ async def test_moderations_response_carries_litellm_call_id_header(): assert fastapi_response.headers["x-litellm-model-id"] == "mod-deployment-1" +@pytest.mark.asyncio +async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplog): + """LIT-7836: the /v1/moderations error line must carry the litellm_call_id the + client sent, rendered in the message and as a structured log record field.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy._types import ProxyException + + call_id = "moderations-call-7836" + + async def passthrough_add_litellm_data(data, **kwargs): + return data + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": "hi"}') + fake_logging = MagicMock() + fake_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + fake_logging.post_call_failure_hook = AsyncMock() + + verbose_proxy_logger.propagate = True + try: + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + finally: + verbose_proxy_logger.propagate = False + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch): from litellm.proxy.agent_endpoints.agent_registry import ( diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 94ccc2762c5..df18e5c6093 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -160,6 +160,37 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): assert "litellm_metadata" not in captured["optional_params"] +@pytest.mark.asyncio +async def test_proxy_only_error_log_keeps_the_request_litellm_call_id(monkeypatch: pytest.MonkeyPatch): + """LIT-7836: a route that already stamped the caller's litellm_call_id must + keep it when the failure is a proxy-only error, so the spend-log row and the + error line share one id instead of a fresh uuid minted here.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + call_id: Final = "caller-supplied-7836" + captured: dict[str, object] = {} + + def fake_pre_call(self, *args, **kwargs): + captured["litellm_call_id"] = self.litellm_call_id + + async def _noop_async_failure(self, *args, **kwargs): + return None + + monkeypatch.setattr(Logging, "pre_call", fake_pre_call) + monkeypatch.setattr(Logging, "async_failure_handler", _noop_async_failure) + request_data: Final[dict[str, object]] = {"model": "gpt-4o", "input": "hi", "litellm_call_id": call_id} + + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/moderations"), + route="/v1/moderations", + original_exception=Exception("bad key"), + ) + + assert request_data["litellm_call_id"] == call_id + assert captured["litellm_call_id"] == call_id + + def test_get_model_group_info_order(): from litellm import Router from litellm.proxy.proxy_server import _get_model_group_info diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index 117c5aa3081..278d11f95b0 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -176,6 +176,21 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): } +@pytest.mark.parametrize( + "exc", + [HTTPException(status_code=401, detail="bad key"), ValueError("provider boom")], + ids=["http_exception", "generic_exception"], +) +def test_handle_exception_on_proxy_returns_the_litellm_call_id_header(exc: Exception): + result = handle_exception_on_proxy(exc, "call-7836") + + assert result.headers == {"x-litellm-call-id": "call-7836"} + + +def test_handle_exception_on_proxy_sends_no_call_id_header_when_the_request_has_none(): + assert handle_exception_on_proxy(ValueError("provider boom")).headers == {} + + @pytest.mark.asyncio async def test_handle_exception_on_proxy_read_only_transaction_forces_writer_recreate( monkeypatch: pytest.MonkeyPatch, From f12feed9a9f0f7008050b3a969e7eeaeec6744e0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 02:39:39 +0000 Subject: [PATCH 15/89] test(proxy): expect litellm_call_id in the image generation call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/proxy_unit_tests/test_proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 9c8dd90dd2b..1fcdaa67143 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -809,6 +809,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth): n=1, size="1024x1024", imageConfig={"aspectRatio": "9:16", "imageSize": "1K"}, + litellm_call_id=mock.ANY, metadata=mock.ANY, proxy_server_request=mock.ANY, secret_fields=mock.ANY, From 1aa2e19ee4dc46a861c65ba7ebb506d41c307908 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 19:52:00 -0700 Subject: [PATCH 16/89] test(together_ai): stop depending on a serverless model we do not control Together moved openai/gpt-oss-20b off serverless, so three tests started failing with a 400 model_not_available from the live API. None of them was really testing Together: they cover provider-prefix parsing, prompt shaping and streaming, all litellm side. Mock the transport and assert those, so the tests answer to our code instead of a vendor catalog. --- tests/local_testing/test_completion.py | 109 ++++++++++++++++++------- 1 file changed, 80 insertions(+), 29 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 43ed57f63af..a6ac112a3b7 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -11,6 +11,7 @@ import io from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm @@ -57,21 +58,49 @@ def test_response_model_none(): assert isinstance(x, litellm.ModelResponse) +TOGETHER_AI_CHAT_URL = "https://api.together.ai/v1/chat/completions" + + +def _together_ai_chat_response(content="Hello!"): + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-oss-20b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), + ) + + def test_completion_custom_provider_model_name(): - try: - litellm.cache = None + litellm.cache = None + with patch.object( + HTTPHandler, "post", return_value=_together_ai_chat_response() + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, logger_fn=logger_fn, + api_key="fake-key", ) - # Add assertions here to check the-response - print(response) - print(response["choices"][0]["finish_reason"]) - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") + + assert mock_post.call_args.kwargs["url"] == TOGETHER_AI_CHAT_URL + assert json.loads(mock_post.call_args.kwargs["data"])["model"] == "openai/gpt-oss-20b" + assert response.choices[0].finish_reason == "stop" def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse: @@ -2804,12 +2833,11 @@ def test_completion_together_ai_llama(): # test_completion_together_ai() def test_customprompt_together_ai(): - try: - litellm.set_verbose = False - litellm.num_retries = 0 - print("in test_customprompt_together_ai") - print(litellm.success_callback) - print(litellm._async_success_callback) + litellm.set_verbose = False + litellm.num_retries = 0 + with patch.object( + HTTPHandler, "post", return_value=_together_ai_chat_response() + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, @@ -2827,14 +2855,14 @@ def test_customprompt_together_ai(): "post_message": "<|im_end|>", }, }, + api_key="fake-key", ) - print(response) - except litellm.exceptions.Timeout as e: - print(f"Timeout Error") - pass - except Exception as e: - print(f"ERROR TYPE {type(e)}") - pytest.fail(f"Error occurred: {e}") + + body = json.loads(mock_post.call_args.kwargs["data"]) + assert body["messages"] == messages + assert "prompt" not in body + assert "roles" not in body + assert response.choices[0].finish_reason == "stop" # test_customprompt_together_ai() @@ -3648,19 +3676,42 @@ def test_completion_together_ai_stream(): litellm.set_verbose = True user_message = "Write 1pg about YC & litellm" messages = [{"content": user_message, "role": "user"}] - try: + sse_body = ( + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"role":"assistant",' + '"content":"YC"},"finish_reason":null}]}\n\n' + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"content":" and ' + 'litellm"},"finish_reason":null}]}\n\n' + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{},' + '"finish_reason":"stop"}]}\n\n' + "data: [DONE]\n\n" + ) + stream_response = httpx.Response( + 200, + content=sse_body.encode(), + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), + ) + + with patch.object( + HTTPHandler, "post", return_value=stream_response + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, stream=True, max_tokens=5, + api_key="fake-key", ) - print(response) - for chunk in response: - print(chunk) - # print(string_response) - except Exception as e: - pytest.fail(f"Error occurred: {e}") + chunks = list(response) + + assert json.loads(mock_post.call_args.kwargs["data"])["stream"] is True + assert "".join( + chunk.choices[0].delta.content or "" for chunk in chunks + ) == "YC and litellm" + assert chunks[-1].choices[0].finish_reason == "stop" # test_completion_together_ai_stream() From be506936bd9b64776a2730b8e4c4c55f4f000491 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:52:40 -0700 Subject: [PATCH 17/89] fix(mcp): preserve explicit caller authorization credentials --- .../_experimental/mcp_server/upstream.py | 11 +++---- .../mcp_server/test_mcp_server_manager.py | 31 +++++++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 49aca5cf31b..2a9312e17a1 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -36,17 +36,16 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True -def validate_static_credential( - server: MCPServer, headers: Mapping[str, str], *, header_slot: str | None = None, openapi: bool = False -) -> Result[None, CredError]: +def validate_static_credential(server: MCPServer, headers: Mapping[str, str]) -> Result[None, CredError]: if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: return Ok(None) default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" slots: Final = frozenset( name.lower() for name in ( - header_slot or server.upstream_token_header or default_slot, - "Authorization" if openapi else default_slot, + server.upstream_token_header or default_slot, + default_slot, + "Authorization", ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) @@ -75,7 +74,7 @@ def validate_openapi_credentials( headers: Final = merge_openapi_headers( server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers ) - match validate_static_credential(server, headers, openapi=True): + match validate_static_credential(server, headers): case Error(error): raise_public(error) case Ok(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index c8dba1c1554..a582632259e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13693,3 +13693,34 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_slot", [None, "X-Custom"]) + @pytest.mark.parametrize("source", ["caller", "forwarded"]) + async def test_api_key_preserves_explicit_authorization_credential( + self, custom_slot: str | None, source: str + ) -> None: + server: Final = MCPServer( + server_id="caller-auth", name="caller-auth", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header=custom_slot, + ) + headers: Final = {"Authorization": "Bearer caller-credential", "X-API-Key": ""} + client: Final = await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=headers if source == "caller" else None, + extra_headers=headers if source == "forwarded" else None, + ) + request: Final = await client.prepare_request_auth() + assert request.headers["Authorization"] == "Bearer caller-credential" + assert request.headers["X-API-Key"] == "" + assert custom_slot is None or custom_slot not in request.headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["", " ", "Bearer", "Basic", "token", "ApiKey"]) + async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: + server: Final = MCPServer( + server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) + assert exc.value.status_code == 500 From 581c613f6680b1daf0e1da140250987d61ede4ac Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 03:02:50 +0000 Subject: [PATCH 18/89] fix(proxy): keep litellm_call_id on shaped errors and list_batches failure hook Already shaped ProxyException and HTTPException errors passing through the moderations, audio speech, Anthropic Messages, and handle_exception_on_proxy paths now answer with the x-litellm-call-id header the route logged under, without overwriting a header the exception was raised with. The GET /v1/batches failure hook receives the resolved request data so the spend log request_id matches the response header and the error log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 5 +- litellm/proxy/batches_endpoints/endpoints.py | 2 +- .../common_utils/openai_error_payload.py | 12 +++- litellm/proxy/proxy_server.py | 18 +++-- litellm/proxy/utils.py | 8 ++- .../anthropic_endpoints/test_endpoints.py | 30 ++++++++ .../proxy/batches_endpoints/test_endpoints.py | 18 +++++ .../common_utils/test_openai_error_payload.py | 25 +++++++ tests/test_litellm/proxy/test_proxy_server.py | 71 ++++++++++++++++++- .../proxy/utils/helpers/test_error_helpers.py | 8 ++- 10 files changed, 184 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f673a82654d..644778bcb9f 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_param, openai_error_type, + with_litellm_call_id, ) from litellm.types.utils import TokenCountResponse @@ -222,7 +223,9 @@ async def anthropic_response( log_llm_api_exception(e, base_llm_response_processor.litellm_call_id) if isinstance(e, ProxyException): - return _anthropic_error_json_response(e, request) + return _anthropic_error_json_response( + with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request + ) # Extract model_id from request metadata (same as success path) litellm_metadata: Final = data.get("litellm_metadata", {}) or {} diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e3767d06e7d..f37c06aea97 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -861,7 +861,7 @@ async def list_batches( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data={"after": after, "limit": limit}, + request_data={**data, "after": after, "limit": limit}, ) litellm_call_id: Final = request_litellm_call_id(data) log_llm_api_exception(e, litellm_call_id) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index d4312d93559..cbc8c78d4f9 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -9,6 +9,9 @@ from typing import Final from fastapi import status from litellm.constants import STRINGIFIED_NONE +from litellm.proxy._types import ProxyException + +LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id" _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { @@ -57,4 +60,11 @@ def openai_error_param(exc: object) -> str | None: def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers if litellm_call_id is None: return None - return {"x-litellm-call-id": litellm_call_id} # mutable-ok: ProxyException mutates its headers dict + return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict + + +def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException: + """The same error object, answering with ``x-litellm-call-id`` when it was raised without one.""" + if litellm_call_id is not None: + exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) + return exc diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index de42dcec8d0..03a425fedb2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -390,7 +390,11 @@ from litellm.proxy.common_utils.model_listing_utils import ( from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) -from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + litellm_call_id_headers, + with_litellm_call_id, +) from litellm.proxy.common_utils.periodic_reload_schedule import ( MODEL_COST_MAP_RELOAD_PARAM_NAME, clear_reload_interval, @@ -11530,7 +11534,7 @@ async def moderations( ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): - raise + raise with_litellm_call_id(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), @@ -11678,8 +11682,14 @@ async def audio_speech( request_data=data, ) log_llm_api_exception(e, litellm_call_id) - if isinstance(e, (ProxyException, HTTPException)): - raise e + if isinstance(e, ProxyException): + raise with_litellm_call_id(e, litellm_call_id) + if isinstance(e, HTTPException): + raise HTTPException( + status_code=e.status_code, + detail=e.detail, + headers={LITELLM_CALL_ID_HEADER: litellm_call_id, **(e.headers or {})}, + ) raise ProxyException( message=getattr(e, "message", f"{e}"), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 53782227998..1376020a907 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,7 +38,11 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) -from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers, openai_error_param +from litellm.proxy.common_utils.openai_error_payload import ( + litellm_call_id_headers, + openai_error_param, + with_litellm_call_id, +) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -7660,7 +7664,7 @@ def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): - return e + return with_litellm_call_id(e, litellm_call_id) _status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) return ProxyException( message=str(e), diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index e4b15cfdcd3..9a9ccd9a213 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -336,6 +336,36 @@ class TestErrorLogCarriesCallId: assert record.litellm_call_id == call_id assert call_id in record.getMessage() + @pytest.mark.asyncio + async def test_messages_already_shaped_failure_answers_with_the_call_id(self): + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + + call_id = "messages-call-7836-shaped" + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_call_id": call_id} + raise ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402) + + request = MagicMock() + request.headers = {} + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + ): + mock_logging.post_call_failure_hook = AsyncMock() + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.status_code == 402 + assert response.headers["x-litellm-call-id"] == call_id + @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): from fastapi import HTTPException diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index cf805b384d2..d9bfb3fe3da 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1976,6 +1976,24 @@ async def test_list__exception_calls_failure_hook(list_harness): assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" +@pytest.mark.asyncio +async def test_list__failure_hook_and_response_share_the_request_litellm_call_id(list_harness): + call_id = "lit7836-list-batches-call-id" + list_harness.pre_call.side_effect = lambda **kw: ( + {**list_harness.body["body"], "litellm_call_id": call_id}, + MagicMock(), + ) + list_harness.litellm_alist.side_effect = ValueError("provider boom") + + with pytest.raises(ProxyException) as raised: + await call_list(list_harness, after="batch-0", limit=5) + + failure_request_data = list_harness.logging.post_call_failure_hook.call_args.kwargs["request_data"] + assert failure_request_data["litellm_call_id"] == call_id + assert (failure_request_data["after"], failure_request_data["limit"]) == ("batch-0", 5) + assert raised.value.headers["x-litellm-call-id"] == call_id + + # =========================================================================== # # # # POST /v1/batches/{batch_id}/cancel - cancel_batch routing-contract tests # diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index df775916046..c09b8742b50 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -9,6 +9,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( litellm_call_id_headers, openai_error_param, openai_error_type, + with_litellm_call_id, ) @@ -164,3 +165,27 @@ def test_a_stringified_none_type_or_param_is_treated_as_absent(): def test_a_failed_request_answers_with_the_call_id_it_was_logged_under(): assert litellm_call_id_headers("call-7836") == {"x-litellm-call-id": "call-7836"} assert litellm_call_id_headers(None) is None + + +def test_an_already_shaped_proxy_error_answers_with_the_call_id_it_was_logged_under(): + raised_without_id = ProxyException(message="budget exceeded", type="budget_exceeded", param="key", code=402) + + carried = with_litellm_call_id(raised_without_id, "call-7836") + + assert carried is raised_without_id + assert carried.headers == {"x-litellm-call-id": "call-7836"} + assert (carried.message, carried.type, carried.param, carried.code) == ( + "budget exceeded", + "budget_exceeded", + "key", + "402", + ) + + +def test_a_proxy_error_keeps_the_call_id_it_was_raised_with(): + raised_with_id = ProxyException( + message="nope", type="None", param=None, code=400, headers={"x-litellm-call-id": "first"} + ) + + assert with_litellm_call_id(raised_with_id, "second").headers == {"x-litellm-call-id": "first"} + assert with_litellm_call_id(ProxyException(message="nope", type="None", param=None, code=400), None).headers == {} diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 85c1a23c473..f3c477e4ad0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -20,7 +20,7 @@ import fastapi.routing import httpx import pytest import yaml -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.encoders import jsonable_encoder from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient @@ -32,7 +32,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyErrorTypes, ProxyException, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, initialize @@ -12970,6 +12970,73 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo assert call_id in record.getMessage() +@pytest.mark.asyncio +async def test_moderations_already_shaped_failure_answers_with_the_callers_litellm_call_id(): + """LIT-7836: a ProxyException raised inside /v1/moderations is re-raised unwrapped but still + answers with the caller's x-litellm-call-id so the client can join it to the error log.""" + call_id = "moderations-call-7836-shaped" + exc = ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402) + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": "hi"}') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + assert raised.value is exc + assert raised.value.code == "402" + assert raised.value.headers["x-litellm-call-id"] == call_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + HTTPException(status_code=401, detail="bad key"), + ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402), + ], + ids=["http_exception", "proxy_exception"], +) +async def test_audio_speech_already_shaped_failure_answers_with_the_callers_litellm_call_id(exc: Exception): + """LIT-7836: /v1/audio/speech re-raises HTTP and proxy shaped failures unchanged, and they must + still answer with the caller's x-litellm-call-id.""" + call_id = "speech-call-7836-shaped" + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"model": "tts-1", "input": "hi", "voice": "alloy"}') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(type(exc)) as raised, + ): + await proxy_server_module.audio_speech( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + if isinstance(exc, HTTPException): + assert (raised.value.status_code, raised.value.detail) == (401, "bad key") + else: + assert raised.value is exc + assert raised.value.headers["x-litellm-call-id"] == call_id + + @pytest.mark.asyncio async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch): from litellm.proxy.agent_endpoints.agent_registry import ( diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index 278d11f95b0..df399c8b1d2 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -178,8 +178,12 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): @pytest.mark.parametrize( "exc", - [HTTPException(status_code=401, detail="bad key"), ValueError("provider boom")], - ids=["http_exception", "generic_exception"], + [ + HTTPException(status_code=401, detail="bad key"), + ValueError("provider boom"), + ProxyException(message="already wrapped", type=ProxyErrorTypes.budget_exceeded.value, param="key", code=402), + ], + ids=["http_exception", "generic_exception", "already_proxy_exception"], ) def test_handle_exception_on_proxy_returns_the_litellm_call_id_header(exc: Exception): result = handle_exception_on_proxy(exc, "call-7836") From 8c046e13bdcbd61f570f10f674ff57b2dcf19afb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 20:17:42 -0700 Subject: [PATCH 19/89] test(together_ai): move request-shape checks to the mapped file, drop the live ones Together moved openai/gpt-oss-20b off serverless and three tests in test_completion.py died on a live 400. None of them needed Together to be up: streaming is already covered live by tests/e2e/llm_translation/test_together_ai_e2e.py, which picks its model from the cost map instead of pinning one, and the other two are request-shape questions. Delete all three and assert the two shapes in the mapped transformation file: the provider prefix is stripped without eating the rest of a slashed model name, and custom role wrappers never reach the request. --- tests/local_testing/test_completion.py | 125 ------------------ .../test_together_ai_chat_transformation.py | 64 +++++++++ 2 files changed, 64 insertions(+), 125 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index a6ac112a3b7..25c6c50251d 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -11,7 +11,6 @@ import io from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest import litellm @@ -58,51 +57,6 @@ def test_response_model_none(): assert isinstance(x, litellm.ModelResponse) -TOGETHER_AI_CHAT_URL = "https://api.together.ai/v1/chat/completions" - - -def _together_ai_chat_response(content="Hello!"): - return httpx.Response( - 200, - json={ - "id": "chatcmpl-together", - "object": "chat.completion", - "created": 1, - "model": "openai/gpt-oss-20b", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - }, - request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), - ) - - -def test_completion_custom_provider_model_name(): - litellm.cache = None - with patch.object( - HTTPHandler, "post", return_value=_together_ai_chat_response() - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - logger_fn=logger_fn, - api_key="fake-key", - ) - - assert mock_post.call_args.kwargs["url"] == TOGETHER_AI_CHAT_URL - assert json.loads(mock_post.call_args.kwargs["data"])["model"] == "openai/gpt-oss-20b" - assert response.choices[0].finish_reason == "stop" - - def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse: new_response = MagicMock() new_response.headers = {"hello": "world"} @@ -2832,40 +2786,6 @@ def test_completion_together_ai_llama(): # test_completion_together_ai() -def test_customprompt_together_ai(): - litellm.set_verbose = False - litellm.num_retries = 0 - with patch.object( - HTTPHandler, "post", return_value=_together_ai_chat_response() - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - roles={ - "system": { - "pre_message": "<|im_start|>system\n", - "post_message": "<|im_end|>", - }, - "assistant": { - "pre_message": "<|im_start|>assistant\n", - "post_message": "<|im_end|>", - }, - "user": { - "pre_message": "<|im_start|>user\n", - "post_message": "<|im_end|>", - }, - }, - api_key="fake-key", - ) - - body = json.loads(mock_post.call_args.kwargs["data"]) - assert body["messages"] == messages - assert "prompt" not in body - assert "roles" not in body - assert response.choices[0].finish_reason == "stop" - - -# test_customprompt_together_ai() def response_format_tests(response: litellm.ModelResponse): @@ -3672,51 +3592,6 @@ async def test_acompletion_stream_watsonx(): # test_maritalk() -def test_completion_together_ai_stream(): - litellm.set_verbose = True - user_message = "Write 1pg about YC & litellm" - messages = [{"content": user_message, "role": "user"}] - sse_body = ( - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"role":"assistant",' - '"content":"YC"},"finish_reason":null}]}\n\n' - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"content":" and ' - 'litellm"},"finish_reason":null}]}\n\n' - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{},' - '"finish_reason":"stop"}]}\n\n' - "data: [DONE]\n\n" - ) - stream_response = httpx.Response( - 200, - content=sse_body.encode(), - headers={"content-type": "text/event-stream"}, - request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), - ) - - with patch.object( - HTTPHandler, "post", return_value=stream_response - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - stream=True, - max_tokens=5, - api_key="fake-key", - ) - chunks = list(response) - - assert json.loads(mock_post.call_args.kwargs["data"])["stream"] is True - assert "".join( - chunk.choices[0].delta.content or "" for chunk in chunks - ) == "YC and litellm" - assert chunks[-1].choices[0].finish_reason == "stop" - - -# test_completion_together_ai_stream() - - def test_moderation(): response = litellm.moderation(input="i'm ishaan cto of litellm") print(response) diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index 7eb7dc41d4f..a7347edb2c7 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -1108,3 +1108,67 @@ def test_get_optional_params_preserves_max_for_declared_levels_model(): ) assert optional_params["reasoning_effort"] == "max" + + +def _together_chat_transport() -> tuple[HTTPHandler, list[httpx.Request]]: + captured_requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1234567890, + "model": TOOL_CALLING_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + return client, captured_requests + + +def test_only_the_provider_prefix_is_stripped_from_a_slashed_model_name(): + client, captured_requests = _together_chat_transport() + + litellm.completion( + model=f"together_ai/{TOOL_CALLING_MODEL}", + messages=[{"role": "user", "content": "Hello!"}], + api_key="fake-key", + client=client, + ) + + assert "/" in TOOL_CALLING_MODEL + assert str(captured_requests[0].url) == "https://api.together.ai/v1/chat/completions" + assert json.loads(captured_requests[0].content)["model"] == TOOL_CALLING_MODEL + + +def test_custom_role_wrappers_never_reach_the_request(): + client, captured_requests = _together_chat_transport() + messages = [{"role": "user", "content": "Hello!"}] + + litellm.completion( + model=f"together_ai/{TOOL_CALLING_MODEL}", + messages=messages, + roles={ + "system": {"pre_message": "<|im_start|>system\n", "post_message": "<|im_end|>"}, + "assistant": {"pre_message": "<|im_start|>assistant\n", "post_message": "<|im_end|>"}, + "user": {"pre_message": "<|im_start|>user\n", "post_message": "<|im_end|>"}, + }, + api_key="fake-key", + client=client, + ) + + request_body = json.loads(captured_requests[0].content) + assert request_body["messages"] == messages + assert "prompt" not in request_body + assert "roles" not in request_body From 258176de762939aca13375abef5a474b4595e8d9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:32:25 -0700 Subject: [PATCH 20/89] fix(mcp): validate rendered static credential payloads --- .../_experimental/mcp_server/upstream.py | 8 ++- .../mcp_server/test_mcp_server_manager.py | 65 ++++++++++++++++++- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 2a9312e17a1..66840db21ce 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -4,7 +4,7 @@ import base64 from collections.abc import Mapping from typing import Final -from litellm.experimental_mcp_client.client import MCPClient +from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError @@ -24,13 +24,17 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False + if auth_type in (MCPAuth.bearer_token, MCPAuth.token): + scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" + credential: Final = strip_auth_scheme(value, scheme).strip() + return bool(credential) and credential.lower() != scheme.lower() if auth_type == MCPAuth.basic: parts: Final = value.split(None, 1) if len(parts) != 2 or parts[0].lower() != "basic": return False try: decoded: Final = base64.b64decode(parts[1], validate=True).strip() - return bool(decoded) and decoded.lower() != b"basic" + return b":" in decoded except ValueError: return False return True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index a582632259e..54add273c24 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13639,7 +13639,7 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM="]) + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM=", "Basic bm8tY29sb24="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", transport=MCPTransport.http, auth_type=MCPAuth.basic) @@ -13724,3 +13724,66 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["no-colon", "Basic bm8tY29sb24="]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_basic_requires_a_username_password_separator(self, value: str, source: str) -> None: + server: Final = MCPServer( + server_id="basic-pair", name="basic-pair", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["user:pass", "user:", ":pass", ":"]) + async def test_basic_preserves_username_password_pairs(self, value: str) -> None: + import base64 + + server: Final = MCPServer( + server_id="basic-valid", name="basic-valid", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value, + ) + client: Final = await MCPServerManager()._create_mcp_client(server) + request: Final = await client.prepare_request_auth() + scheme, encoded = request.headers["Authorization"].split(" ", 1) + assert scheme == "Basic" + assert base64.b64decode(encoded) == value.encode() + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value", [ + (MCPAuth.bearer_token, "Bearer"), (MCPAuth.bearer_token, "Bearer "), (MCPAuth.bearer_token, "bearer"), + (MCPAuth.token, "token"), (MCPAuth.token, "token "), (MCPAuth.token, "TOKEN"), + ]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_static_scheme_only_input_cannot_hide_behind_rendered_prefix( + self, auth_type: MCPAuthType, value: str, source: str + ) -> None: + server: Final = MCPServer( + server_id="empty-scheme", name="empty-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + authentication_token=value if source == "configured" else None, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value,expected", [ + (MCPAuth.bearer_token, "token", "Bearer token"), + (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), + (MCPAuth.token, "tokenish", "token tokenish"), + ]) + async def test_static_credentials_that_resemble_schemes_remain_usable( + self, auth_type: MCPAuthType, value: str, expected: str + ) -> None: + server: Final = MCPServer( + server_id="real-token", name="real-token", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value, + ) + client: Final = await MCPServerManager()._create_mcp_client(server) + request: Final = await client.prepare_request_auth() + assert request.headers["Authorization"] == expected From 2b6184d76867fd38a990d2df124b7d1cd808ca6c Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:08:24 +0000 Subject: [PATCH 21/89] fix(proxy): resolve rate-limit fallbacks after model normalization and retry from a client-request snapshot The fallback retry in _pre_call_with_fallbacks re-entered common_processing_pre_call_logic with data already enriched by the first pass, so add_litellm_data_to_request deep-copied a metadata dict holding the live OTel span and the request failed with a 500 (cannot pickle '_thread.RLock') instead of the intended 429 or fallback. Capture the configured fallbacks and a snapshot of the client request before the first pass, look up the fallback chain by the normalized model group after the limiter raises, and run each fallback attempt on a fresh copy of that snapshot. Replaces the mock-heavy tests with a rig that runs the real v3 limiter and a live OTel span through the proxy_logging_obj seam Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 62 ++-- .../test_response_polling_pre_call_checks.py | 8 +- .../proxy/test_common_request_processing.py | 347 +++++++----------- 3 files changed, 166 insertions(+), 251 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3c8bb9b3c92..7c9a296965c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2066,20 +2066,12 @@ class ProxyBaseLLMRequestProcessing: ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - original_model: Final = self.data.get("model") - fallback_models: Final = ( - self._resolve_fallback_models( - model=original_model, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - ) - if original_model - and isinstance(original_model, str) - and llm_router - and not self.data.get("disable_fallbacks") + configured_fallbacks: Final = ( + self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict) + if llm_router is not None and not self.data.get("disable_fallbacks") else None ) - pristine: Final = independent_snapshot(self.data) if fallback_models else None + pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None try: return await self.common_processing_pre_call_logic( @@ -2099,7 +2091,16 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) except ProxyRateLimitError as original_exc: - if not fallback_models or pristine is None: + rate_limited_data: Final = self.data + original_model: Final = rate_limited_data.get("model") + if pristine is None or not configured_fallbacks or not isinstance(original_model, str): + raise + + fallback_models: Final = self._resolve_fallback_models( + model=original_model, + fallbacks=configured_fallbacks, + ) + if not fallback_models: raise verbose_proxy_logger.info( @@ -2133,39 +2134,30 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError: continue except BaseException: - self.data = pristine + self.data = rate_limited_data raise - self.data = pristine + self.data = rate_limited_data raise original_exc - def _resolve_fallback_models( - self, - model: str, - llm_router: Router, - user_api_key_dict: UserAPIKeyAuth, - ) -> list | None: - from litellm.router_utils.fallback_event_handlers import get_fallback_model_group - - fallbacks = None - + @staticmethod + def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None: key_router_settings: Final = user_api_key_dict.router_settings - if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: - fallbacks = key_router_settings["fallbacks"] + key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None + fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks + return fallbacks if isinstance(fallbacks, list) and fallbacks else None - if fallbacks is None: - fallbacks = llm_router.fallbacks - - if not fallbacks: - return None + @staticmethod + def _resolve_fallback_models(model: str, fallbacks: list) -> list | None: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group fallback_model_group, generic_fallback_idx = get_fallback_model_group( fallbacks=fallbacks, model_group=model, ) - if fallback_model_group is None and generic_fallback_idx is not None: - fallback_model_group = fallbacks[generic_fallback_idx]["*"] - return fallback_model_group + if fallback_model_group is not None: + return fallback_model_group + return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None @staticmethod def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str: diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 38f087f51ca..459834d0fd2 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -48,10 +48,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(fallbacks=None), + llm_router=MagicMock(), general_settings={}, proxy_config=MagicMock(), skip_pre_call_logic=True, @@ -87,10 +87,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(fallbacks=None), + llm_router=MagicMock(), general_settings={}, proxy_config=MagicMock(), ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 92af65b2637..f689dd62df6 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6365,247 +6365,170 @@ class TestPreCallWithFallbacksOnLocalRateLimit: call_type="acompletion", ) - @pytest.mark.asyncio - async def test_fallback_retries_from_pristine_request_data(self): - import threading + @staticmethod + def _v3_limiter_rig( + monkeypatch: pytest.MonkeyPatch, + user_api_key_dict: ProxyUserAPIKeyAuth, + fallbacks: list[dict[str, list[str]]], + ) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]: + """Real v3 limiter (the default ``parallel_request_limiter``) wired in through the + ``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real: + ``add_litellm_data_to_request`` with a live OTel span, ``function_setup``, then the limiter.""" + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + monkeypatch.setattr(proxy_server, "prisma_client", None) + limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + limiter_models: list[str] = [] - primary_model = "gpt-4" - fallback_model = "gpt-3.5-turbo" + async def run_limiter(**kwargs): + limiter_models.append(kwargs["data"]["model"]) + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=kwargs["data"], + call_type=kwargs["call_type"], + ) + return kwargs["data"] - processor = ProxyBaseLLMRequestProcessing( - data={ - "model": primary_model, - "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, - } + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) + router = litellm.Router( + model_list=[ + {"model_name": group, "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}} + for chain in fallbacks + for group in (*chain.keys(), *(m for models in chain.values() for m in models)) + ], + fallbacks=fallbacks, ) + return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models - metadata_at_entry = [] - - async def mock_pre_call_logic(**kwargs): - copy.deepcopy(processor.data["metadata"]) - metadata_at_entry.append(dict(processor.data["metadata"])) - processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() - processor.data["litellm_logging_obj"] = object() - if processor.data.get("model") == primary_model: - raise ProxyRateLimitError( - detail="TPM limit exceeded for gpt-4", - headers={"retry-after": "30"}, - ) - return processor.data, MagicMock() - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: [fallback_model]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - data, logging_obj = await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model=primary_model, - route_type="acompletion", - llm_router=mock_router, - ) - - assert processor.data["model"] == fallback_model - assert metadata_at_entry[1] == {"tags": ["a"]} - - @pytest.mark.asyncio - async def test_exhausted_fallbacks_restore_pristine_request_data(self): - import threading - - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - - primary_model = "gpt-4" - original_data = { - "model": primary_model, - "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, - } - processor = ProxyBaseLLMRequestProcessing(data=copy.deepcopy(original_data)) - - async def mock_pre_call_logic(**kwargs): - processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() - processor.data["litellm_logging_obj"] = object() - raise ProxyRateLimitError( - detail=f"TPM limit exceeded for {processor.data.get('model')}", - headers={"retry-after": "30"}, - ) - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: ["gpt-3.5-turbo"]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - with pytest.raises(ProxyRateLimitError, match="gpt-4"): - await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model=primary_model, - route_type="acompletion", - llm_router=mock_router, - ) - - assert processor.data == original_data - - @pytest.mark.asyncio - async def test_real_add_litellm_data_to_request_rerun_with_otel_span_falls_back(self): - from opentelemetry import trace + @staticmethod + def _otel_key(**limits) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - from litellm.proxy.proxy_server import ProxyConfig + span = TracerProvider().get_tracer("test").start_span("proxy-request") + return ProxyUserAPIKeyAuth(api_key="hashed-key", parent_otel_span=span, **limits) - trace.set_tracer_provider(TracerProvider()) + @staticmethod + def _chat_request() -> Request: + return Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}) - primary_model = "gpt-4" - fallback_model = "gpt-3.5-turbo" - - request_mock = MagicMock(spec=Request) - request_mock.url = MagicMock() - request_mock.url.path = "/v1/chat/completions" - request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" - request_mock.method = "POST" - request_mock.query_params = {} - request_mock.headers = {"Content-Type": "application/json"} - request_mock.client = MagicMock() - request_mock.client.host = "127.0.0.1" - - user_api_key_dict = UserAPIKeyAuth( - parent_otel_span=trace.get_tracer("x").start_span("s"), - api_key="hashed-key", - user_id="u1", - team_id="t1", - metadata={}, - team_metadata={}, - team_member_tpm_limit=1000, + async def _pre_call( + self, + data: dict, + user_api_key_dict: ProxyUserAPIKeyAuth, + rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], + ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict, object]]: + proxy_logging_obj, router, proxy_config, _ = rig + processor = ProxyBaseLLMRequestProcessing(data=data) + result = await processor._pre_call_with_fallbacks( + request=self._chat_request(), + general_settings={}, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + version=None, + proxy_config=proxy_config, + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=None, + route_type="acompletion", + llm_router=router, ) + return processor, result - processor = ProxyBaseLLMRequestProcessing( - data={ + @pytest.mark.asyncio + async def test_v3_limiter_with_otel_span_falls_back_from_client_request(self, monkeypatch: pytest.MonkeyPatch): + """Customer path: OTel on, per-key model RPM cap on the primary, a router fallback configured. + The first pass enriches ``data["metadata"]`` with the live span, then the limiter raises. The + fallback pass must start from the client's request again, so ``add_litellm_data_to_request`` + never deep-copies the span (the ``cannot pickle '_thread.RLock'`` 500).""" + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + + def client_request() -> dict: + return { "model": primary_model, "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, + "metadata": {"tags": ["client-tag"]}, } - ) - async def real_add_litellm_data_pre_call(**kwargs): - await add_litellm_data_to_request( - data=processor.data, - request=request_mock, - user_api_key_dict=user_api_key_dict, - proxy_config=ProxyConfig(), + _, (first_data, _) = await self._pre_call(client_request(), key, rig) + processor, (data, logging_obj) = await self._pre_call(client_request(), key, rig) + + assert first_data["model"] == primary_model + assert data["model"] == fallback_model + assert processor.data is data + assert data["litellm_logging_obj"] is logging_obj + assert logging_obj.model == fallback_model + requester_metadata = data["metadata"]["requester_metadata"] + assert requester_metadata["tags"] == ["client-tag"] + assert "litellm_parent_otel_span" not in requester_metadata + assert "user_api_key_auth" not in requester_metadata + assert data["metadata"]["litellm_parent_otel_span"] is key.parent_otel_span + assert rig[3] == [primary_model, primary_model, fallback_model] + + @pytest.mark.asyncio + async def test_v3_limiter_with_otel_span_returns_429_when_fallbacks_exhausted( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(rpm_limit=1) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig) + processor = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await processor._pre_call_with_fallbacks( + request=self._chat_request(), general_settings={}, - version="test", - ) - if processor.data.get("model") == primary_model: - raise ProxyRateLimitError( - detail="TPM limit exceeded for gpt-4", - headers={"retry-after": "30"}, - ) - return processor.data, MagicMock() - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: [fallback_model]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=real_add_litellm_data_pre_call, - ): - data, logging_obj = await processor._pre_call_with_fallbacks( - request=request_mock, - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=user_api_key_dict, + proxy_logging_obj=rig[0], + user_api_key_dict=key, version=None, - proxy_config=MagicMock(), + proxy_config=rig[2], user_model=None, user_temperature=None, user_request_timeout=None, user_max_tokens=None, user_api_base=None, - model=primary_model, + model=None, route_type="acompletion", - llm_router=mock_router, + llm_router=rig[1], ) - assert processor.data["model"] == fallback_model + assert rig[3] == [primary_model, primary_model, fallback_model] + assert exc_info.value.status_code == 429 + assert "Rate limit exceeded" in str(exc_info.value.detail) + assert exc_info.value.headers["retry-after"] + assert processor.data["model"] == primary_model + assert processor.data["litellm_logging_obj"].model == primary_model + assert processor.data["litellm_call_id"] @pytest.mark.asyncio - async def test_no_fallbacks_skips_snapshot(self): - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + async def test_fallback_lookup_uses_alias_resolved_model_group(self, monkeypatch: pytest.MonkeyPatch): + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + monkeypatch.setattr(litellm, "model_alias_map", {"my-alias": primary_model}) + key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": "my-alias", "messages": [{"role": "user", "content": "hi"}]} - processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) - async def mock_pre_call_logic(**kwargs): - raise ProxyRateLimitError( - detail="TPM limit exceeded", - headers={"retry-after": "30"}, - ) - - mock_router = MagicMock() - mock_router.fallbacks = None - - with patch( # test-quality-ok: spying the snapshot seam is the only observable check that the no-fallback path skips it - "litellm.proxy.common_request_processing.independent_snapshot" - ) as snapshot_mock: - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - with pytest.raises(ProxyRateLimitError): - await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model="gpt-4", - route_type="acompletion", - llm_router=mock_router, - ) - - snapshot_mock.assert_not_called() + assert data["model"] == fallback_model + assert rig[3] == [primary_model, primary_model, fallback_model] class _RecordingSuccessLogger(CustomLogger): From 9faaf7f4d436b558dbc4d1ab25a5801269629d6d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:15:13 +0000 Subject: [PATCH 22/89] refactor(proxy): assign the fallback model on the fresh snapshot instead of building a dict literal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7c9a296965c..d18fab1c9f0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2113,7 +2113,8 @@ class ProxyBaseLLMRequestProcessing: for fallback_model in fallback_models: if fallback_model == original_model: continue - self.data = {**independent_snapshot(pristine), "model": fallback_model} + self.data = independent_snapshot(pristine) + self.data["model"] = fallback_model try: return await self.common_processing_pre_call_logic( request=request, From 391da46e2cdcef46491519ba2814b5d38c752285 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:59:33 +0000 Subject: [PATCH 23/89] test: type the v3 limiter rig and otel key helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_common_request_processing.py | 34 ++++++++++++------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f689dd62df6..7d9a78fd981 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -44,6 +44,7 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.utils import ProxyLogging @@ -6383,15 +6384,17 @@ class TestPreCallWithFallbacksOnLocalRateLimit: limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) limiter_models: list[str] = [] - async def run_limiter(**kwargs): - limiter_models.append(kwargs["data"]["model"]) + async def run_limiter( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: + limiter_models.append(str(data["model"])) await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=DualCache(), - data=kwargs["data"], - call_type=kwargs["call_type"], + data=data, + call_type=call_type, ) - return kwargs["data"] + return data proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) @@ -6406,11 +6409,18 @@ class TestPreCallWithFallbacksOnLocalRateLimit: return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models @staticmethod - def _otel_key(**limits) -> ProxyUserAPIKeyAuth: + def _otel_key( + rpm_limit: int | None = None, model_rpm_limit: dict[str, int] | None = None + ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider span = TracerProvider().get_tracer("test").start_span("proxy-request") - return ProxyUserAPIKeyAuth(api_key="hashed-key", parent_otel_span=span, **limits) + return ProxyUserAPIKeyAuth( + api_key="hashed-key", + parent_otel_span=span, + rpm_limit=rpm_limit, + metadata={"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}, + ) @staticmethod def _chat_request() -> Request: @@ -6418,10 +6428,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: async def _pre_call( self, - data: dict, + data: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], - ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict, object]]: + ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict[str, object], LiteLLMLoggingObj]]: proxy_logging_obj, router, proxy_config, _ = rig processor = ProxyBaseLLMRequestProcessing(data=data) result = await processor._pre_call_with_fallbacks( @@ -6450,10 +6460,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: never deep-copies the span (the ``cannot pickle '_thread.RLock'`` 500).""" primary_model = "gpt-4.1" fallback_model = "gpt-4.1-mini" - key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + key = self._otel_key(model_rpm_limit={primary_model: 1}) rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) - def client_request() -> dict: + def client_request() -> dict[str, object]: return { "model": primary_model, "messages": [{"role": "user", "content": "hi"}], @@ -6520,7 +6530,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: primary_model = "gpt-4.1" fallback_model = "gpt-4.1-mini" monkeypatch.setattr(litellm, "model_alias_map", {"my-alias": primary_model}) - key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + key = self._otel_key(model_rpm_limit={primary_model: 1}) rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) request = {"model": "my-alias", "messages": [{"role": "user", "content": "hi"}]} From 9974cf4bf817331053649fccc1f1b9e00d61d570 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 08:17:58 +0000 Subject: [PATCH 24/89] fix(proxy): honor key-level disable_fallbacks after first pre-call pass Key metadata disable_fallbacks only lands on data during add_key_level_controls, so the local rate-limit fallback retry now rechecks it post pre-call. Also use a real UserAPIKeyAuth in the skip pre-call test since the path reads router_settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 7 ++++- .../test_response_polling_pre_call_checks.py | 2 +- .../proxy/test_common_request_processing.py | 30 +++++++++++++++++-- 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d18fab1c9f0..36b49ef064c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2093,7 +2093,12 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data original_model: Final = rate_limited_data.get("model") - if pristine is None or not configured_fallbacks or not isinstance(original_model, str): + if ( + pristine is None + or not configured_fallbacks + or rate_limited_data.get("disable_fallbacks") + or not isinstance(original_model, str) + ): raise fallback_models: Final = self._resolve_fallback_models( diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..9f1a228855e 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -87,7 +87,7 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=UserAPIKeyAuth(), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, llm_router=MagicMock(), diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7d9a78fd981..e9d74a46daa 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6410,7 +6410,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @staticmethod def _otel_key( - rpm_limit: int | None = None, model_rpm_limit: dict[str, int] | None = None + rpm_limit: int | None = None, + model_rpm_limit: dict[str, int] | None = None, + disable_fallbacks: bool = False, ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider @@ -6419,7 +6421,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: api_key="hashed-key", parent_otel_span=span, rpm_limit=rpm_limit, - metadata={"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}, + metadata={ + **({"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}), + **({"disable_fallbacks": True} if disable_fallbacks else {}), + }, ) @staticmethod @@ -6540,6 +6545,27 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert data["model"] == fallback_model assert rig[3] == [primary_model, primary_model, fallback_model] + @pytest.mark.asyncio + async def test_key_metadata_disable_fallbacks_returns_429_instead_of_retrying( + self, monkeypatch: pytest.MonkeyPatch + ): + """``disable_fallbacks`` set in key metadata only lands on ``data`` during the first + pre-call pass (``add_key_level_controls``), so it must be honored after that pass.""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(model_rpm_limit={primary_model: 1}, disable_fallbacks=True) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._pre_call(dict(request), key, rig) + + assert exc_info.value.status_code == 429 + assert rig[3] == [primary_model, primary_model] + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From 2d40254b57a62be9bce95586cb902b90c8505c4c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:15:46 -0700 Subject: [PATCH 25/89] feat(e2e): key the provider cache per test and mount Bedrock behind it The exact-request cache reused 5% of routed traffic (build 218: 19 hits, 350 misses) because every test salts its prompt with a fresh unique_marker(), so the same test could never match itself across builds. It also routed only openai and anthropic, while the week's flakiness was Bedrock. Key is now HMAC(test id + method + URL + headers + body, with every unique_marker() token replaced by a placeholder, + FIFO slot index). The slot index is what keeps two marker-only-different calls in one test on two recordings and therefore two provider response ids, so spend rows still reconcile one per invocation. A call outside any test is not cacheable. Bedrock gets a region-qualified mount and SigV4 re-signing, since the edge rewrites the Host the proxy signed. Signature headers are excluded from the key for signing mounts only, because x-amz-date would otherwise make every Bedrock request a permanent miss; every other mount still keys on its credentials whole. Only Anthropic-on-Bedrock chat deployments route: embeddings, image generation, rerank and realtime keep their direct path, and so do deployments carrying their own aws_role_name or static keys, whose whole point is to prove the product's assume-role chain rather than the runner's. The two eventstream actions bypass the cache and go live, still signed. Counters are now attributed per mount as well as in total, so a build can report a per-provider hit rate instead of one number. --- .../test_provider_cache.py | 509 +++++++++++++++--- tests/e2e/fixture_canonical.py | 5 +- tests/e2e/models.py | 1 + tests/e2e/provider_cache.py | 171 ++++-- tests/e2e/provider_cache_routing.py | 47 +- tests/e2e/provider_edge.py | 62 ++- tests/e2e/provider_edge_bedrock.py | 72 +++ tests/e2e/test_provider_edge.py | 19 +- 8 files changed, 768 insertions(+), 118 deletions(-) create mode 100644 tests/e2e/provider_edge_bedrock.py diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 828227ed239..5d35981344d 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -7,7 +7,7 @@ import subprocess import threading import time import uuid -from collections.abc import Generator +from collections.abc import Generator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace @@ -18,21 +18,58 @@ from urllib.parse import urlsplit import pytest from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward -from models import LiteLLMParamsBody -from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response +from models import LiteLLMParamsBody, ModelMode +from botocore.credentials import Credentials +from provider_cache import ( + SIGNATURE_HEADERS, + CacheEdge, + CacheHit, + CaptureLease, + ResponseStore, + cacheable_endpoint, + request_identity, + slotted_key, + successful_response, +) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model -from provider_edge import configured_cache_backend, start_provider_edge +from fixture_mode import SESSION_TEST_KEY +from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge +from provider_edge_bedrock import bedrock_signer from redis.exceptions import ConnectionError as RedisConnectionError SECRET: Final = b"synthetic-cache-hmac-key-for-tests" BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}' SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"} +TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_case" +OTHER_TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_other_case" + + +def marked(marker: str) -> bytes: + """One request body shaped like the suite's own: a fixed prompt salted with a + 12-lowercase-hex ``unique_marker()`` token, fresh on every run.""" + return b'{"model":"test","messages":[{"role":"user","content":"hello %s"}]}' % marker.encode() + + +MARKED: Final = marked("0a1b2c3d4e5f") +BEDROCK_MOUNT: Final = "bedrock/us-east-1" +BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1%3A0" +BEDROCK_BODY: Final = b'{"messages":[{"role":"user","content":[{"text":"hello 0a1b2c3d4e5f"}]}]}' +CONVERSE_SUCCESS: Final = ( + b'{"output":{"message":{"role":"assistant","content":[{"text":"hi"}]}},' + b'"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1,"totalTokens":2}}' +) +INVOKE_SUCCESS: Final = ( + b'{"id":"msg_synthetic","type":"message","role":"assistant",' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}' +) +STATIC_CREDENTIALS: Final = Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") class Provider(ThreadingHTTPServer): hits: tuple[tuple[str, bytes], ...] = () + authorizations: tuple[str, ...] = () response: bytes = SUCCESS status: int = 200 delay: float = 0 @@ -49,6 +86,7 @@ class Handler(BaseHTTPRequestHandler): assert isinstance(server, Provider) body: Final = self.rfile.read(int(self.headers.get("content-length", "0"))) server.hits += ((self.path, body),) + server.authorizations += (self.headers.get("authorization", ""),) time.sleep(server.delay) self.send_response(server.status) if server.stream: @@ -122,6 +160,29 @@ def store(redis_url: str) -> RedisResponseStore: return redis_store(redis_url, "test-" + uuid.uuid4().hex) +def cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + """A cache edge standing in for one pytest process. A fresh instance over the + same store is the next build running the same test: the recordings survive, + the per-test FIFO slot counters start over.""" + return CacheEdge(store, SECRET, test_key=lambda: test_key) + + +def slot_key( + url: str, slot: int = 0, body: bytes | None = BODY, + headers: dict[str, str] = HEADERS, test_key: str = TEST_KEY, +) -> str: + prepared: Final = prepare_forward("POST", url, headers, body) + assert isinstance(prepared, PreparedForward) + return slotted_key(SECRET, request_identity(SECRET, test_key, "POST", url, prepared.headers, body), slot) + + +def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + return CacheEdge( + store, SECRET, test_key=lambda: test_key, + signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + ) + + @contextmanager def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: upstream: Final = f"http://127.0.0.1:{provider.server_port}" @@ -132,34 +193,54 @@ def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: running.shutdown() +@contextmanager +def bedrock_edge(cache: CacheEdge, provider: Provider, action: str = "converse") -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={BEDROCK_MOUNT: upstream}) + try: + yield f"{running.edge.api_base(BEDROCK_MOUNT)}/model/{BEDROCK_MODEL}/{action}" + finally: + running.shutdown() + + def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse: result: Final = forward("POST", url, headers=headers, body=body, timeout=5) assert isinstance(result, RawResponse), result return result -def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: +def test_repeated_call_takes_its_own_slot_and_both_replay_next_run( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS - with edge(CacheEdge(store, SECRET), provider) as other: + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as other: assert call(other).body == SUCCESS - assert len(provider.hits) == 1 + assert call(other).body == SUCCESS + assert len(provider.hits) == 2 @pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")]) def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, body) + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: call(url, body) assert len(provider.hits) == 2 @pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")]) def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, headers=HEADERS | {name: value}) call(url + "?x=1") assert len(provider.hits) == 3 @@ -169,36 +250,59 @@ def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provid def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None: provider.status = status provider.response = response - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).status_code == status + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: assert call(url).body == response assert len(provider.hits) == 2 def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None: provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure" - with edge(CacheEdge(store, SECRET), provider) as url: - replies: Final = tuple(call(url) for _ in range(2)) + with edge(cache_edge(store), provider) as url: + live: Final = call(url) + with edge(cache_edge(store), provider) as url: + replayed: Final = call(url) assert len(provider.hits) == 1 - assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies) + assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in (live, replayed)) def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None: short: Final = replace(store, lifetime_ms=250) - with edge(CacheEdge(short, SECRET), provider) as url: - call(url) - call(url) - time.sleep(0.3) - call(url) - call(url) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + + def drain() -> None: + head = cache_edge(short).forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + + drain() + assert len(provider.hits) == 1 + drain() + assert len(provider.hits) == 1 + time.sleep(0.3) + drain() assert len(provider.hits) == 2 -def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None: +def test_concurrent_builds_publish_one_recording_atomically( + store: RedisResponseStore, provider: Provider, +) -> None: + """Five processes running the same test at the same time all reach slot 0 of + one key, which is the only way the capture lease is contended now that a + repeat inside a single test takes its own slot.""" provider.delay = 0.15 - with edge(CacheEdge(store, SECRET), provider) as url: - with ThreadPoolExecutor(max_workers=5) as executor: - replies: Final = tuple(executor.map(lambda _: call(url).body, range(5))) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + edges: Final = tuple(cache_edge(store) for _ in range(5)) + + def drain(cache: CacheEdge) -> bytes: + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + return b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) + + with ThreadPoolExecutor(max_workers=5) as executor: + replies: Final = tuple(executor.map(drain, edges)) assert replies == (SUCCESS,) * 5 assert len(provider.hits) == 1 @@ -231,9 +335,9 @@ def test_stream_completion_controls_publication(store: RedisResponseStore, provi provider.stream = True provider.truncated = truncated provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' - with edge(CacheEdge(store, SECRET), provider) as url: - for _ in range(2): - result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) + for _ in range(2): + with edge(cache_edge(store), provider) as url: + result = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) if truncated: assert isinstance(result, NetworkError) else: @@ -246,9 +350,9 @@ def test_store_outage_preserves_provider_success(provider: Provider) -> None: probe.bind(("127.0.0.1", 0)) port: Final = probe.getsockname()[1] unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable") - with edge(CacheEdge(unavailable, SECRET), provider) as url: - assert call(url).body == SUCCESS - assert call(url).body == SUCCESS + for _ in range(2): + with edge(cache_edge(unavailable), provider) as url: + assert call(url).body == SUCCESS assert len(provider.hits) == 2 @@ -267,7 +371,8 @@ def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None def test_identity_preserves_values_and_never_contains_credentials() -> None: variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'') - keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants) + url: Final = "https://example.invalid/v1/chat/completions" + keys: Final = tuple(request_identity(SECRET, TEST_KEY, "POST", url, HEADERS, body) for body in variants) assert len(set(keys)) == len(variants) assert all(len(key) == 64 and "synthetic-account" not in key for key in keys) @@ -275,21 +380,22 @@ def test_identity_preserves_values_and_never_contains_credentials() -> None: @pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()]) def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY) + key: Final = slot_key(upstream) lease: Final = store.lookup(key) assert isinstance(lease, CaptureLease) assert store.publish(key, lease, payload) - cache: Final = CacheEdge(store, SECRET) - for _ in range(2): - head = cache.forward("POST", upstream, HEADERS, BODY, 5) + caches: Final = tuple(cache_edge(store) for _ in range(2)) + for cache in caches: + head = cache.forward("openai", "POST", upstream, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 1 - assert dict(cache.counters.counts) == { - "corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1, + assert dict(caches[0].counters.counts) == { + "corrupt": 1, "mount:openai:corrupt": 1, "misses": 1, "mount:openai:misses": 1, + "upstream_attempts": 1, "mount:openai:upstream_attempts": 1, + "writes": 1, "mount:openai:writes": 1, } + assert dict(caches[1].counters.counts) == {"hits": 1, "mount:openai:hits": 1} @pytest.mark.parametrize("payload", [ @@ -301,22 +407,242 @@ def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisRespon def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: provider.stream = True provider.response = payload - with edge(CacheEdge(store, SECRET), provider) as url: - assert call(url).body == payload - assert call(url).body == payload + for _ in range(2): + with edge(cache_edge(store), provider) as url: + assert call(url).body == payload assert len(provider.hits) == 2 +def test_requests_differing_only_by_marker_share_one_recording_per_slot( + store: RedisResponseStore, provider: Provider, +) -> None: + """The whole point of the canonical key. Every e2e test salts its prompt with + a fresh ``unique_marker()``, so before this the same test could never reuse + anything across builds. The second run mints markers it has never sent, which + is what a later build actually does, and must still serve both from the two + slots the first run recorded.""" + with edge(cache_edge(store), provider) as url: + assert call(url, MARKED).body == SUCCESS + assert call(url, marked("f5e4d3c2b1a0")).body == SUCCESS + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: + assert call(url, marked("7c6b5a493827")).body == SUCCESS + assert call(url, marked("1122334455ff")).body == SUCCESS + assert len(provider.hits) == 2 + + +@pytest.mark.parametrize("body", [ + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5f0"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0A1B2C3D4E5F"}]}', + b'{"model":"0a1b2c3d4e5f","messages":[{"role":"user","content":"hello"}]}', +]) +def test_a_token_that_is_not_a_marker_keeps_its_own_key( + store: RedisResponseStore, provider: Provider, body: bytes, +) -> None: + """Too short, too long, upper case, or in another field: none of these is the + 12-lowercase-hex token ``unique_marker`` mints, so none may fold onto it.""" + with edge(cache_edge(store), provider) as url: + call(url, MARKED) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: + call(url, body) + assert len(provider.hits) == 2 + + +def test_another_test_never_reuses_this_tests_recording( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: + call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + + +def test_calls_outside_any_test_are_never_cached( + store: RedisResponseStore, provider: Provider, +) -> None: + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + cache: Final = CacheEdge(store, SECRET, test_key=lambda: SESSION_TEST_KEY) + for _ in range(2): + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts) == { + "bypass": 2, "mount:openai:bypass": 2, + "upstream_attempts": 2, "mount:openai:upstream_attempts": 2, + } + + +def test_counters_attribute_every_outcome_to_its_mount( + store: RedisResponseStore, provider: Provider, +) -> None: + """The build report needs per-provider hit counts, and the flat totals cannot + supply them. Anthropic is served a chat-shaped body here, which its validator + rejects, so one mount writes and the other does not.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cache: Final = cache_edge(store) + running: Final = start_provider_edge(cache, mounts={"openai": upstream, "anthropic": upstream}) + try: + call(running.edge.api_base("openai") + "/v1/chat/completions") + call(running.edge.api_base("anthropic") + "/v1/messages") + finally: + running.shutdown() + counts: Final = dict(cache.counters.counts) + assert counts["misses"] == 2 + assert counts["mount:openai:misses"] == 1 and counts["mount:anthropic:misses"] == 1 + assert counts["mount:openai:writes"] == 1 and "mount:anthropic:writes" not in counts + assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts + + +class TestBedrockSigning: + """Bedrock is the reason the edge could not mount it before: SigV4 covers the + Host header, so forwarding through a rewritten api_base invalidates the + proxy's signature. The edge mints its own over the upstream URL instead.""" + + def test_the_proxys_signature_is_replaced_not_forwarded(self) -> None: + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + signed: Final = signer( + "POST", + f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse", + {"content-type": "application/json", "Authorization": "AWS4-HMAC-SHA256 Credential=PROXY/...", + "X-Amz-Date": "19700101T000000Z", "X-Amz-Security-Token": "proxy-session-token"}, + BEDROCK_BODY, + ) + assert "PROXY" not in str(signed) and "proxy-session-token" not in str(signed) + assert signed["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + assert "/us-east-1/bedrock/aws4_request" in signed["Authorization"] + assert signed["X-Amz-Date"] != "19700101T000000Z" + assert signed["content-type"] == "application/json" + + def test_the_signed_url_reaches_the_wire_byte_for_byte(self) -> None: + """SigV4 hashes the canonical URI, so if the HTTP layer re-encoded the + colon in an inference-profile id after signing, every call would fail + with a signature mismatch rather than anything that names the cause.""" + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse" + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + prepared: Final = prepare_forward("POST", url, signer("POST", url, dict(HEADERS), BEDROCK_BODY), BEDROCK_BODY) + assert isinstance(prepared, PreparedForward) + assert urlsplit(prepared.url).path == urlsplit(url).path + + def test_signature_headers_are_excluded_from_the_key( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A real signature is fresh on every call, so keying on it would make + every Bedrock request a permanent miss. The stub signer here varies its + stamp per call on purpose: the real one only varies once a second, which + would let this pass by luck when it should fail.""" + provider.response = CONVERSE_SUCCESS + stamps: Final = iter(("20260101T000000Z", "20260102T111111Z")) + + def varying(method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} + + def signing_edge() -> CacheEdge: + return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + + for _ in range(2): + with bedrock_edge(signing_edge(), provider) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 1 + assert provider.authorizations[0] == ( + f"AWS4-HMAC-SHA256 http://127.0.0.1:{provider.server_port}/model/{BEDROCK_MODEL}/converse" + ), "the signature must cover the upstream URL the edge calls, not the edge URL the proxy called" + + def test_a_mount_without_a_signer_still_keys_on_its_credentials( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """The exclusion is per mount. Dropping authorization globally would let + one OpenAI account read another's recording.""" + cache: Final = bedrock_cache_edge(store) + assert "authorization" in SIGNATURE_HEADERS + assert "authorization" in cache.keyed("openai", HEADERS) + assert "authorization" not in cache.keyed(BEDROCK_MOUNT, HEADERS) + with edge(cache, provider) as url: + call(url) + with edge(bedrock_cache_edge(store), provider) as url: + call(url, headers=HEADERS | {"authorization": "Bearer synthetic-account-two"}) + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action,response", [("converse", CONVERSE_SUCCESS), ("invoke", INVOKE_SUCCESS)]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("action,response", [ + ("converse", b'{"output":{"message":{}}}'), + ("converse", b'{"stopReason":"end_turn"}'), + ("converse", b'{"message":"The provided model identifier is invalid."}'), + ("converse", CONVERSE_SUCCESS[:-20]), + ("invoke", b'{"id":"msg_x","type":"message","content":[{"type":"text","text":"hi"}]}'), + ("invoke", b'{"id":"msg_x","type":"message","stop_reason":"end_turn"}'), + ("invoke", b'{"message":"Too many requests, please wait before trying again."}'), + ]) + def test_incomplete_or_error_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) + def test_streaming_endpoints_go_live_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, + ) -> None: + """An eventstream's completeness cannot be proven without parsing its + frames, so these bypass rather than risk recording a truncated answer. + They are still signed: a bypass is a forward, not a passthrough.""" + provider.response = CONVERSE_SUCCESS + cache: Final = bedrock_cache_edge(store) + for _ in range(2): + with bedrock_edge(cache, provider, action) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + assert all( + sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + for sent in provider.authorizations + ), provider.authorizations + + @pytest.mark.parametrize("action,cacheable", [ + ("converse", True), ("invoke", True), + ("converse-stream", False), ("invoke-with-response-stream", False), + ]) + def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + + def test_a_region_mount_resolves_whole(self) -> None: + resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) + assert resolved is not None + assert resolved.mount == BEDROCK_MOUNT + assert resolved.upstream_base == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert resolved.upstream_path == f"model/{BEDROCK_MODEL}/converse" + + def test_anthropic_stream_requires_start_finish_and_stop() -> None: start: Final = b'data: {"type":"message_start","message":{}}\n\n' finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n' stop: Final = b'data: {"type":"message_stop"}\n\n' url: Final = "https://example.invalid/v1/messages" headers: Final = {"content-type": "text/event-stream"} - assert successful_response(url, 200, headers, start + finish + stop) - assert not successful_response(url, 200, headers, start + stop) - assert not successful_response(url, 200, headers, finish + stop) - assert not successful_response(url, 200, headers, start + finish) + assert successful_response("anthropic", url, 200, headers, start + finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + stop) + assert not successful_response("anthropic", url, 200, headers, finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + finish) @pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) @@ -342,6 +668,53 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params +@pytest.mark.parametrize("model", [ + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/converse/us.anthropic.claude-sonnet-5", + "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +]) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: + params: Final = LiteLLMParamsBody(model=model) + routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) + assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" + assert routed.api_base is None + assert routed.model_dump(exclude={"aws_bedrock_runtime_endpoint"}) == params.model_dump( + exclude={"aws_bedrock_runtime_endpoint"} + ) + + +@pytest.mark.parametrize("params", [ + LiteLLMParamsBody(model="bedrock/amazon.titan-embed-text-v2:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-canvas-v1:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-sonic-v1:0"), + LiteLLMParamsBody(model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_role_name="arn:aws:iam::1:role/x"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_access_key_id="AKIA"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", api_base="https://custom.invalid"), + LiteLLMParamsBody( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_bedrock_runtime_endpoint="https://custom.invalid", + ), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), +]) +def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: + """Non-Anthropic models the runner role cannot invoke, deployments carrying + their own AWS identity (routing those would replace the assume-role chain the + batch suite exists to prove), explicit endpoints, and unmounted regions.""" + routed: Final = route_cache_model( + params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, + ) + assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + + +@pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) +def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: + params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + assert route_cache_model( + params, lambda mount: f"http://edge.invalid/{mount}", enabled=True, mode=mode, + ) is params + + def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None: params: Final = LiteLLMParamsBody(model="openai/test") assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params @@ -366,14 +739,16 @@ class PublishOutage: def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None: unavailable: Final = replace(store, client=PublishOutage(store.client)) - cache: Final = CacheEdge(unavailable, SECRET) + cache: Final = cache_edge(unavailable) with edge(cache, provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS assert len(provider.hits) == 2 assert dict(cache.counters.counts)["write_failures"] == 2 - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS + assert len(provider.hits) == 3 + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert len(provider.hits) == 3 @@ -382,45 +757,41 @@ def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> with socket.socket() as unavailable: unavailable.bind(("127.0.0.1", 0)) url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError) - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + cache: Final = cache_edge(store) + assert isinstance(cache.forward("openai", "POST", url, dict(HEADERS), BODY, 0.2), NetworkError) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) assert dict(cache.counters.counts)["rejected"] == 1 def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - head: Final = cache.forward("POST", url, HEADERS, BODY, 5) + cache: Final = cache_edge(store) + head: Final = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) head.steps.close() - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) def test_effective_account_change_cannot_reuse_cache( store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - for account in ("account-a", "account-b", "account-b"): + caches: Final = tuple(cache_edge(store) for _ in range(3)) + for account, cache in zip(("account-a", "account-b", "account-b"), caches, strict=True): netrc = tmp_path / account netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n") monkeypatch.setenv("NETRC", str(netrc)) - head = cache.forward("POST", url, HEADERS, BODY, 5) + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 2 - assert dict(cache.counters.counts)["hits"] == 1 + assert dict(caches[2].counters.counts)["hits"] == 1 def test_enabled_environment_reuses_store_across_fresh_backends( @@ -449,7 +820,7 @@ def test_enabled_environment_reuses_store_across_fresh_backends( def test_duplicate_headers_bypass_cache_and_count_live_calls( store: RedisResponseStore, provider: Provider, known_mount: bool, ) -> None: - cache: Final = CacheEdge(store, SECRET) + cache: Final = cache_edge(store) with edge(cache, provider) as url: parsed: Final = urlsplit(url) for _ in range(2): diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py index e76d63ca33b..019c011aa67 100644 --- a/tests/e2e/fixture_canonical.py +++ b/tests/e2e/fixture_canonical.py @@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = ( ) SECRET_PLACEHOLDER: Final = "" +MARKER_PATTERN: Final = re.compile(r"(?" + PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( (re.compile(r"(?"), ( @@ -67,7 +70,7 @@ PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"), "", ), - (re.compile(r"(?"), + (MARKER_PATTERN, MARKER_PLACEHOLDER), ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..7550bfdc150 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -951,6 +951,7 @@ class LiteLLMParamsBody(BaseModel): aws_access_key_id: str | None = None aws_secret_access_key: str | None = None aws_region_name: str | None = None + aws_bedrock_runtime_endpoint: str | None = None vertex_project: str | None = None vertex_location: str | None = None vertex_credentials: str | None = None diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0c6eac75a43..1dc2f99abe5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -23,12 +23,18 @@ from e2e_http import ( prepare_forward, primed_steps, ) +from fixture_canonical import MARKER_PATTERN, MARKER_PLACEHOLDER +from fixture_mode import SESSION_TEST_KEY, current_test_key from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError LIFETIME_SECONDS: Final = 86_400 MAX_REQUEST_BYTES: Final = 256 * 1024 MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024 UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"}) +SIGNATURE_HEADERS: Final = frozenset( + {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} +) +BEDROCK_MOUNT_PREFIX: Final = "bedrock" JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -56,6 +62,7 @@ class CacheUnavailable: type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable +type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] class ResponseStore(Protocol): @@ -83,28 +90,51 @@ class SignedResponse(BaseModel): signature: str -def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str: +def canonical_text(value: str) -> str: + return MARKER_PATTERN.sub(MARKER_PLACEHOLDER, value) + + +def canonical_body(body: bytes) -> bytes: + try: + return canonical_text(body.decode("utf-8")).encode("utf-8") + except UnicodeDecodeError: + return body + + +def request_identity( + secret: bytes, test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> str: fields: Final = ( - b"provider-cache-exact-v1", method.encode(), url.encode(), + b"provider-cache-canonical-v2", test_key.encode(), method.encode(), canonical_text(url).encode(), *(part.encode() for pair in sorted(headers.items()) for part in pair), - b"no-body" if body is None else b"body", b"" if body is None else body, + b"no-body" if body is None else b"body", b"" if body is None else canonical_body(body), ) encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields) return hmac.new(secret, encoded, hashlib.sha256).hexdigest() -def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool: - return ( - method == "POST" - and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"} - and body is not None - and len(body) <= MAX_REQUEST_BYTES - ) +def slotted_key(secret: bytes, identity: str, slot: int) -> str: + return hmac.new(secret, f"{identity}:{slot}".encode(), hashlib.sha256).hexdigest() -def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: +def is_bedrock(mount: str) -> bool: + return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX + + +def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: + if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: + return False + path: Final = urlsplit(url).path + if is_bedrock(mount): + return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path in {"/v1/chat/completions", "/v1/messages"} + + +def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES: return False + if is_bedrock(mount): + return complete_bedrock_response(url, body) streaming: Final = "text/event-stream" in headers.get("content-type", "").lower() if streaming: try: @@ -147,6 +177,26 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body: ) +def complete_bedrock_response(url: str, body: bytes) -> bool: + """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an + Anthropic model answers the Anthropic message shape. Either way a truncated + or error body is missing the terminator field, which is what makes it safe to + record. The streaming variants never reach here: they are not cacheable.""" + try: + value: Final = JSON_VALUE.validate_json(body) + except ValidationError: + return False + if not isinstance(value, dict) or "message" in value: + return False + if urlsplit(url).path.endswith("/converse"): + return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) + return ( + value.get("type") == "message" + and isinstance(value.get("content"), list) + and isinstance(value.get("stop_reason"), str) + ) + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -172,7 +222,7 @@ def encode_response(secret: bytes, response: CachedResponse) -> bytes: return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode() -def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None: +def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: str) -> CachedResponse | None: if len(payload) > 2 * MAX_RESPONSE_BYTES: return None try: @@ -183,7 +233,9 @@ def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> Cached chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks) except (ValidationError, ValueError): return None - if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)): + if response.request_key != key or not successful_response( + mount, url, response.status_code, response.headers, b"".join(chunks) + ): return None return response @@ -199,6 +251,24 @@ class CacheCounters: self.counts = tuple((current | {name: current.get(name, 0) + 1}).items()) +@dataclass(slots=True) +class SlotCounter: + """FIFO position of a request among the canonically identical ones its test + has already sent. Two calls in one test that differ only by ``unique_marker`` + canonicalize the same, so without this they would share one recording and the + second would replay the first's provider response id.""" + + counts: tuple[tuple[str, int], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def take(self, identity: str) -> int: + with self.lock: + current: Final = dict(self.counts) + taken: Final = current.get(identity, 0) + self.counts = tuple((current | {identity: taken + 1}).items()) + return taken + + @dataclass(slots=True) class ResponseCapture: buffer: io.BytesIO = field(default_factory=io.BytesIO) @@ -231,9 +301,12 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + slots: SlotCounter = field(default_factory=SlotCounter) + signers: Mapping[str, RequestSigner] = field(default_factory=dict) wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep + test_key: Callable[[], str] = current_test_key def lookup(self, key: str) -> CacheLookup: deadline: Final = self.clock() + self.wait_seconds @@ -241,39 +314,71 @@ class CacheEdge: self.sleep(min(0.05, max(0, deadline - self.clock()))) return result - def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError: - if not cacheable_endpoint(method, url, body): - self.counters.increment("bypass") - self.counters.increment("upstream_attempts") - return forward_stream(method, url, headers=headers, body=body, timeout=timeout) - prepared: Final = prepare_forward(method, url, headers, body) + def count(self, mount: str, name: str) -> None: + self.counters.increment(name) + self.counters.increment(f"mount:{mount}:{name}") + + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: + """The headers actually sent upstream. A signing mount gets a signature + minted over the upstream URL, because the edge rewrote the Host the proxy + signed and Bedrock verifies it.""" + signer: Final = self.signers.get(mount) + return headers if signer is None else signer(method, url, headers, body) + + def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: + """A signing mount's signature headers are the edge's own and carry a + timestamp, so keying on them would make every request a permanent miss. + Every other mount keys on its headers whole, credentials included, so a + different account can never read another's recording.""" + if mount not in self.signers: + return headers + return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + + def forward( + self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, + ) -> StreamHead | NetworkError: + test_key: Final = self.test_key() + if test_key == SESSION_TEST_KEY or not cacheable_endpoint(mount, method, url, body): + self.count(mount, "bypass") + self.count(mount, "upstream_attempts") + return forward_stream( + method, url, headers=self.outbound(mount, method, url, headers, body), body=body, timeout=timeout, + ) + prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.counters.increment("rejected") + self.count(mount, "rejected") return prepared - key: Final = exact_key(self.secret, method, url, prepared.headers, body) + identity: Final = request_identity( + self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, + ) + key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): - response: Final = decode_response(self.secret, key, found.payload, url) + response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: - self.counters.increment("hits") + self.count(mount, "hits") return StreamHead(response.status_code, response.headers, response_steps(response)) - self.counters.increment("corrupt" if response is None else "expired") + self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found - self.counters.increment("misses") + self.count(mount, "misses") if isinstance(capture_slot, CacheUnavailable): - self.counters.increment("cache_errors") - self.counters.increment("upstream_attempts") + self.count(mount, "cache_errors") + self.count(mount, "upstream_attempts") head: Final = forward_prepared_stream(prepared, timeout) if not isinstance(capture_slot, CaptureLease): return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.counters.increment("rejected") + self.count(mount, "rejected") return head - return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head))) + return StreamHead( + head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), + ) - def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]: + def capture( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, + ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() try: with closing(head.steps): @@ -285,15 +390,15 @@ class CacheEdge: headers: Final = { name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS } - if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)): - self.counters.increment("rejected") + if not capture.eligible or not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + self.count(mount, "rejected") return response: Final = CachedResponse( request_key=key, status_code=head.status_code, headers=headers, chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), ) published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) - self.counters.increment("writes" if published else "write_failures") + self.count(mount, "writes" if published else "write_failures") finally: self.store.release(key, lease) capture.buffer.close() diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 24599b5a313..e7e4899eb71 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -8,14 +8,57 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) +DEFAULT_BEDROCK_REGION: Final = "us-east-1" +BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." + + +def bedrock_mount(params: LiteLLMParamsBody) -> str | None: + """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + + Only the Anthropic models route. The edge validates converse and invoke + bodies by their Anthropic and Converse terminator fields, and the runner role + is allowed to invoke exactly those models, so Bedrock embeddings, image + generation, rerank and realtime keep their existing direct path rather than + reaching an edge that could neither sign nor validate for them.""" + route: Final = params.model.partition("/")[2] + model: Final = route.partition("/")[2] or route + if BEDROCK_ANTHROPIC_INFIX not in model: + return None + return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + + +def route_bedrock( + params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None, +) -> LiteLLMParamsBody: + """Deployments that carry their own AWS identity stay off the edge. The edge + re-signs with the run pod's role, so routing an `aws_role_name` deployment + would quietly replace the very assume-role chain that test exists to prove.""" + if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None: + return params + if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None: + return params + mount: Final = bedrock_mount(params) + if mount is None: + return params + base: Final = base_for(mount) + if base is None: + return params + return params.model_copy(update={"aws_bedrock_runtime_endpoint": base}) + def route_cache_model( params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None, ) -> LiteLLMParamsBody: - if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None: + if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None: + return params + if params.litellm_credential_name is not None: return params provider: Final = params.model.partition("/")[0] - if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None: + if provider == "bedrock": + return route_bedrock(params, base_for, mode) + if mode == "realtime" or params.api_base is not None: + return params + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index dda9e6f8e4f..8f718ad1967 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -48,7 +48,7 @@ import threading from collections import deque from collections.abc import Generator, Mapping, Sequence from contextlib import closing, contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path @@ -94,17 +94,41 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge +from provider_cache import CacheEdge, RequestSigner, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter +BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",) + EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + **{ + f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" + for region in BEDROCK_REGIONS + }, } ) + +@dataclass(frozen=True, slots=True) +class ResolvedMount: + mount: str + upstream_base: str + upstream_path: str + + +def resolve_mount(path: str, mounts: Mapping[str, str]) -> ResolvedMount | None: + """Longest mount prefix wins, so a region-qualified mount such as + ``bedrock/us-east-1`` resolves whole instead of leaving the region as the + first segment of the upstream path.""" + trimmed: Final = path.lstrip("/") + for mount in sorted(mounts, key=len, reverse=True): + if trimmed == mount or trimmed.startswith(f"{mount}/"): + return ResolvedMount(mount, mounts[mount], trimmed[len(mount):].lstrip("/")) + return None + REPLAY_MISS_STATUS: Final = 599 _HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset( @@ -754,14 +778,14 @@ def _handle_record( def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, - cache: CacheEdge | None = None, + cache: CacheEdge | None = None, mount: str = "", ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } head: Final = ( forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) - if cache is None else cache.forward(method, url, forwarded, body, timeout) + if cache is None else cache.forward(mount, method, url, forwarded, body, timeout) ) match head: case NetworkError(message=message): @@ -796,10 +820,13 @@ def handle_edge_request( prefix, then record (forward + persist) or replay (serve from the bundle). Socket-free so unit tests exercise every branch without a server.""" split: Final = urlsplit(raw_path) - mount, _, upstream_path = split.path.lstrip("/").partition("/") - upstream_base: Final = mounts.get(mount) - if upstream_base is None: - return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}") + resolved: Final = resolve_mount(split.path, mounts) + if resolved is None: + unknown: Final = split.path.lstrip("/").partition("/")[0] + return _text_reply(404, f"unknown provider mount {unknown!r}; known mounts: {', '.join(sorted(mounts))}") + mount: Final = resolved.mount + upstream_base: Final = resolved.upstream_base + upstream_path: Final = resolved.upstream_path profile: Final = ( backend.recorder.profile if isinstance(backend, RecordEdge) @@ -830,7 +857,8 @@ def handle_edge_request( match backend: case CacheEdge(): return _handle_live( - method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, + method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, + backend, mount, ) case LiveEdge(): return _handle_live( @@ -891,7 +919,7 @@ class _EdgeHandler(BaseHTTPRequestHandler): ) if isinstance(edge_server.backend, CacheEdge) and duplicate_headers: edge_server.backend.counters.increment("duplicate_header_bypass") - if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts: + if resolve_mount(urlsplit(self.path).path, edge_server.mounts) is not None: edge_server.backend.counters.increment("upstream_attempts") outcome: Final = handle_edge_request( selected_backend, @@ -1079,6 +1107,8 @@ def provider_edge_api_base( return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount) return None case "record" | "replay": + if is_bedrock(mount): + return None if mount not in EDGE_MOUNTS: raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}") return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base( @@ -1108,7 +1138,17 @@ def configured_cache_backend() -> CacheEdge | None: return None from provider_cache_redis import configured_cache - return configured_cache() + cache: Final = configured_cache() + return None if cache is None else replace(cache, signers=bedrock_signers()) + + +@functools.lru_cache(maxsize=1) +def bedrock_signers() -> Mapping[str, RequestSigner]: + """One signer per mounted Bedrock region, built lazily so a run that never + mounts Bedrock neither imports botocore nor resolves an AWS identity.""" + from provider_edge_bedrock import bedrock_signer + + return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py new file mode 100644 index 00000000000..5d8148482d7 --- /dev/null +++ b/tests/e2e/provider_edge_bedrock.py @@ -0,0 +1,72 @@ +"""SigV4 re-signing for Bedrock traffic routed through the provider edge. + +Bedrock is the one provider the edge could never mount. SigV4 signs the Host +header, so rewriting ``api_base`` to point at the edge invalidates the proxy's +signature and Bedrock rejects the call before it reaches a model. The edge +therefore has to drop the proxy's signature and mint its own over the upstream +URL it is actually about to call. + +The identity it signs with is the run pod's own, from the EKS Pod Identity +association on ServiceAccount ``buildkite-e2e-run``. That role carries Bedrock +invoke and converse on an allowlist of the Anthropic models the suite registers +and nothing else, so a re-signed call can reach exactly the models the suite +already uses. The proxy's own Bedrock credentials are not involved in a routed +deployment, which is why ``aws_role_name`` deployments stay off the edge: their +whole point is to prove the product's assume-role chain. + +Signature headers are excluded from the cache key by the caller, and they have +to be: ``x-amz-date`` is a timestamp, so keying on it would make every Bedrock +request a permanent miss. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials +from botocore.session import Session +from provider_cache import SIGNATURE_HEADERS + +BEDROCK_SERVICE: Final = "bedrock" + + +class MissingAwsCredentials(RuntimeError): + """No AWS identity is resolvable, so the edge cannot sign for Bedrock.""" + + +@dataclass(frozen=True, slots=True) +class BedrockSigner: + region: str + credentials: Callable[[], Credentials] + + def __call__(self, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + unsigned: Final = { + name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS + } + request: Final = AWSRequest(method=method, url=url, headers=unsigned, data=body or b"") + SigV4Auth(self.credentials(), BEDROCK_SERVICE, self.region).add_auth(request) + return dict(request.headers) + + +@functools.lru_cache(maxsize=1) +def pod_credentials() -> Credentials: + """The run pod's own identity, resolved once per process through botocore's + ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" + resolved: Final = Session().get_credentials() + if resolved is None: + raise MissingAwsCredentials( + "the provider edge is mounted for Bedrock but no AWS credentials resolve; " + "the run pod gets them from the Pod Identity association on buildkite-e2e-run" + ) + return resolved + + +def bedrock_signer(region: str, credentials: Callable[[], Credentials] = pod_credentials) -> BedrockSigner: + """Credentials are resolved on the first signed request, not here, so a run + that mounts Bedrock but never calls it needs no AWS identity at all.""" + return BedrockSigner(region, credentials) diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 5d0c79f26f6..c8d70697182 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -1279,15 +1279,30 @@ class TestApiBaseSeam: ) def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None: - with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"): + with pytest.raises(ValueError, match="unknown provider mount 'cohere'"): provider_edge_api_base( - "bedrock", + "cohere", mode_raw="record", bundle_dir=tmp_path / "bundle", bind_host="127.0.0.1", advertise_host="127.0.0.1", ) + @pytest.mark.parametrize("mode_raw", ["record", "replay"]) + def test_bedrock_never_wires_a_bundle_because_the_edge_cannot_sign_into_one( + self, tmp_path: Path, mode_raw: str, + ) -> None: + """Record and replay serve from a bundle without re-signing, so a Bedrock + deployment pointed at that edge would send the proxy's signature over a + rewritten Host. It keeps its direct route in both modes.""" + assert provider_edge_api_base( + "bedrock/us-east-1", + mode_raw=mode_raw, + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) is None + def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None: root = tmp_path / "bundle" first = provider_edge_api_base( From b68e60f7061a2102fa076ff7ca6368f819d13770 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:34:09 -0700 Subject: [PATCH 26/89] feat(e2e): cache the responses and embeddings endpoints behind the edge Chat completions and messages were the only cacheable paths. The suite also drives /v1/embeddings and /v1/responses through the same OpenAI mount, so both now cache, each with its own completeness rule: a chat response's `choices` check would reject a perfectly good embedding, and a Responses run that never reached `response.completed` must stay out of the cache the same way a truncated stream does. Vertex and Gemini stay off the edge. litellm's `_check_custom_proxy` rewrites a path-prefixed vertex api_base into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without a root-mounted edge on its own port or a change in litellm. Shipping an unvalidated URL guess would have been worse than saying so in PROVIDER_CACHE.md. Also finishes the MountPolicy move: a mount now carries its signer and its unkeyed headers together instead of a bare signer map. --- .../test_provider_cache.py | 102 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 28 ++++- tests/e2e/provider_cache.py | 63 +++++++++-- tests/e2e/provider_edge.py | 15 ++- tests/e2e/provider_edge_bedrock.py | 2 +- 5 files changed, 186 insertions(+), 24 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5d35981344d..5c491e1cdcd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -25,6 +25,7 @@ from provider_cache import ( CacheEdge, CacheHit, CaptureLease, + MountPolicy, ResponseStore, cacheable_endpoint, request_identity, @@ -179,7 +180,9 @@ def slot_key( def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: return CacheEdge( store, SECRET, test_key=lambda: test_key, - signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + policies={BEDROCK_MOUNT: MountPolicy( + sign=bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS), unkeyed_headers=SIGNATURE_HEADERS, + )}, ) @@ -501,6 +504,98 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +EMBEDDING_SUCCESS: Final = ( + b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' + b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' +) +RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_STREAM_SUCCESS: Final = ( + b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' +) + + +@contextmanager +def openai_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={"openai": upstream}) + try: + yield running.edge.api_base("openai") + path + finally: + running.shutdown() + + +class TestNonChatOpenAiEndpoints: + """Chat and messages were the only cacheable paths. Embeddings and responses + are the other two JSON endpoints the suite drives through the same mount, and + each needs its own completeness rule: a chat response's ``choices`` check + would reject a perfectly good embedding.""" + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", EMBEDDING_SUCCESS), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + def test_a_completed_response_stream_replays( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + provider.stream = True + provider.response = RESPONSE_STREAM_SUCCESS + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == RESPONSE_STREAM_SUCCESS + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", b'{"object":"list","data":[],"usage":{"prompt_tokens":0}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[]}],"usage":{}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1]}]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"incomplete","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"in_progress","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","output":[]}'), + ]) + def test_incomplete_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("payload", [ + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\n', + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\ndata: {"type":"response.failed"}\n\n', + b'data: {"type":"response.completed","response":{"id":"resp_x"}}\n\ndata: {"type":"response.created"}\n\n', + ]) + def test_a_response_stream_that_never_completed_is_never_cached( + self, store: RedisResponseStore, provider: Provider, payload: bytes, + ) -> None: + provider.stream = True + provider.response = payload + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == payload + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + ("/v1/chat/completions", True), ("/v1/messages", True), + ("/v1/embeddings", True), ("/v1/responses", True), + ("/v1/audio/speech", False), ("/v1/images/generations", False), + ("/v1/files", False), ("/v1/batches", False), + ]) + def test_only_the_json_endpoints_are_cacheable(self, path: str, cacheable: bool) -> None: + assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -545,7 +640,10 @@ class TestBedrockSigning: return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} def signing_edge() -> CacheEdge: - return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + return CacheEdge( + store, SECRET, test_key=lambda: TEST_KEY, + policies={BEDROCK_MOUNT: MountPolicy(sign=varying, unkeyed_headers=SIGNATURE_HEADERS)}, + ) for _ in range(2): with bedrock_edge(signing_edge(), provider) as url: diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 8635c9ed9ae..393a96e8c16 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,11 +1,29 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +## Request identity + +A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is + +Requests that differ only by their markers therefore share a canonical identity, which is what makes the cache reusable across builds: every e2e test salts its prompt afresh, so an exact-byte key would miss on every call. Within one test, calls that share a canonical identity are still recorded and replayed separately, by a FIFO slot index appended to the key. That matters because a replayed response carries the recorded provider response id, `LiteLLM_SpendLogs.request_id` is that id, and one shared recording answering two calls would collapse two spend rows into one + +Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to + +Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure +## Bedrock + +Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss + +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate + +Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm + Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires ## Configuration @@ -18,16 +36,16 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay ## Recorded response semantics -Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching +Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Spend reconciliation keeps its distinct-ID and row-count assertions: its prompts differ by an index as well as a marker, so they stay distinct once markers are normalized, and calls that are canonically equal within one test take separate FIFO slots and separate recordings anyway. Accounting tests are not automatically excluded from caching Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers ## Qualification -`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence +`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis, including the marker-canonical key, the FIFO slot index, per-test isolation, SigV4 re-signing against a local upstream, and each endpoint's completeness rule. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 1dc2f99abe5..a2b18a3a466 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -9,6 +9,7 @@ import time from collections.abc import Callable, Generator, Mapping from contextlib import closing from dataclasses import dataclass, field +from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit @@ -35,6 +36,7 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -65,6 +67,23 @@ type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] +@dataclass(frozen=True, slots=True) +class MountPolicy: + """What a mount needs beyond plain forwarding. + + ``sign`` mints a fresh credential over the upstream URL, for providers whose + auth covers the Host the edge rewrote. ``unkeyed_headers`` names headers that + must stay out of the cache key because they change on every call and would + otherwise make the mount a permanent miss: a minted signature, or an OAuth + token the provider rotates. Naming one costs the guarantee that a recording + can never cross credentials, so a mount with a rotating token relies on the + environment holding one identity for that provider. Mounts with a static API + key name nothing here and keep the guarantee whole.""" + + sign: RequestSigner | None = None + unkeyed_headers: frozenset[str] = frozenset() + + class ResponseStore(Protocol): def lookup(self, key: str) -> CacheLookup: ... @@ -127,7 +146,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) - return path in {"/v1/chat/completions", "/v1/messages"} + return path in OPENAI_JSON_PATHS def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: @@ -150,6 +169,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): return False + if urlsplit(url).path == "/v1/responses": + return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) return ( @@ -168,8 +189,17 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or "error" in value: return False - if urlsplit(url).path == "/v1/messages": + path: Final = urlsplit(url).path + if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) + if path == "/v1/embeddings": + data: Final = value.get("data") + return isinstance(data, list) and bool(data) and isinstance(value.get("usage"), dict) and all( + isinstance(item, dict) and isinstance(item.get("embedding"), list) and bool(item["embedding"]) + for item in data + ) + if path == "/v1/responses": + return value.get("object") == "response" and value.get("status") == "completed" choices: Final = value.get("choices") return isinstance(choices, list) and bool(choices) and all( isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str) @@ -197,6 +227,14 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: + """The Responses API streams typed events and ends with ``response.completed``. + A run that failed, was cancelled, or ran out of tokens ends with a different + terminal event, so requiring that one keeps a half-finished response out.""" + last: Final = values[-1] + return isinstance(last, dict) and last.get("type") == "response.completed" + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -296,13 +334,16 @@ def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None yield StreamChunk(base64.b64decode(chunk, validate=True)) +NO_POLICIES: Final[Mapping[str, MountPolicy]] = MappingProxyType({}) + + @dataclass(frozen=True, slots=True) class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) slots: SlotCounter = field(default_factory=SlotCounter) - signers: Mapping[str, RequestSigner] = field(default_factory=dict) + policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep @@ -321,18 +362,18 @@ class CacheEdge: def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy - signed and Bedrock verifies it.""" - signer: Final = self.signers.get(mount) + signed and the provider verifies it.""" + signer: Final = self.policies.get(mount, MountPolicy()).sign return headers if signer is None else signer(method, url, headers, body) def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: - """A signing mount's signature headers are the edge's own and carry a - timestamp, so keying on them would make every request a permanent miss. - Every other mount keys on its headers whole, credentials included, so a - different account can never read another's recording.""" - if mount not in self.signers: + """Headers the cache key is built from. A mount keeps its credentials in + the key unless its policy names them unkeyed, so by default one account + can never read another's recording.""" + unkeyed: Final = self.policies.get(mount, MountPolicy()).unkeyed_headers + if not unkeyed: return headers - return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + return {name: value for name, value in headers.items() if name.lower() not in unkeyed} def forward( self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 8f718ad1967..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -94,7 +94,7 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge, RequestSigner, is_bedrock +from provider_cache import SIGNATURE_HEADERS, CacheEdge, MountPolicy, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter @@ -1139,16 +1139,21 @@ def configured_cache_backend() -> CacheEdge | None: from provider_cache_redis import configured_cache cache: Final = configured_cache() - return None if cache is None else replace(cache, signers=bedrock_signers()) + return None if cache is None else replace(cache, policies=bedrock_policies()) @functools.lru_cache(maxsize=1) -def bedrock_signers() -> Mapping[str, RequestSigner]: - """One signer per mounted Bedrock region, built lazily so a run that never +def bedrock_policies() -> Mapping[str, MountPolicy]: + """One policy per mounted Bedrock region, built lazily so a run that never mounts Bedrock neither imports botocore nor resolves an AWS identity.""" from provider_edge_bedrock import bedrock_signer - return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) + return MappingProxyType( + { + f"bedrock/{region}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS) + for region in BEDROCK_REGIONS + } + ) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py index 5d8148482d7..73e4a16d272 100644 --- a/tests/e2e/provider_edge_bedrock.py +++ b/tests/e2e/provider_edge_bedrock.py @@ -58,7 +58,7 @@ def pod_credentials() -> Credentials: """The run pod's own identity, resolved once per process through botocore's ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" resolved: Final = Session().get_credentials() - if resolved is None: + if resolved is None: # pyright: ignore[reportUnnecessaryComparison] # stubs miss the empty-chain None raise MissingAwsCredentials( "the provider edge is mounted for Bedrock but no AWS credentials resolve; " "the run pod gets them from the Pod Identity association on buildkite-e2e-run" From aebfcf7da3b5b13591232f4b559978d4f92aafb5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:45:06 -0700 Subject: [PATCH 27/89] fix(e2e): route Bedrock deployments whose region only the proxy can resolve Almost every Bedrock deployment in the suite declares aws_region_name="os.environ/AWS_REGION". The mount resolver treated that string as a region name, produced a mount nothing serves, and left the whole Anthropic-on-Bedrock surface on its direct path, which is the one thing mounting Bedrock was for. The run pod does not share the proxy's environment, so the harness genuinely cannot resolve that reference. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. --- .../test_provider_cache.py | 27 +++++++++++++------ tests/e2e/PROVIDER_CACHE.md | 2 ++ tests/e2e/provider_cache_routing.py | 23 +++++++++++++++- 3 files changed, 43 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5c491e1cdcd..0278bf8fd48 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -766,13 +766,20 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params -@pytest.mark.parametrize("model", [ - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/converse/us.anthropic.claude-sonnet-5", - "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +@pytest.mark.parametrize("model,region", [ + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/converse/us.anthropic.claude-sonnet-5", None), + ("bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), + ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), ]) -def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: - params: Final = LiteLLMParamsBody(model=model) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: + """Almost every Bedrock deployment in the suite declares its region as + `os.environ/AWS_REGION`, which only the proxy can resolve. Treating that + string as a region name would leave the whole Anthropic-on-Bedrock surface + off the edge, which is the point of mounting it at all.""" + params: Final = LiteLLMParamsBody(model=model, aws_region_name=region) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" assert routed.api_base is None @@ -794,15 +801,19 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: aws_bedrock_runtime_endpoint="https://custom.invalid", ), LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), + LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying their own AWS identity (routing those would replace the assume-role chain the - batch suite exists to prove), explicit endpoints, and unmounted regions.""" + batch suite exists to prove), explicit endpoints, unmounted regions, and a + region only the proxy can resolve on a model that is not cross-region, whose + real region the harness cannot know.""" routed: Final = route_cache_model( params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, ) - assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + assert routed is params @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 393a96e8c16..c530574cda2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -20,6 +20,8 @@ An eligible miss calls the provider. A complete successful response is stored im Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss +Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. + Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index e7e4899eb71..d2237bb49bc 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -10,6 +10,26 @@ LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_requ DEFAULT_BEDROCK_REGION: Final = "us-east-1" BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." +BEDROCK_CROSS_REGION_PREFIX: Final = "us." +ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def bedrock_region(declared: str | None, model: str) -> str | None: + """The region whose edge mount a deployment belongs to, or None when the + harness cannot know it. + + Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy + resolves from its own environment. The run pod does not share that + environment, so the harness genuinely does not know the region. A `us.` + inference profile fans out across the US regions and is reachable from any + of them, so the default entry point is correct for those whatever the proxy + resolved; anything else keeps its direct path rather than being sent to a + region the model may not exist in.""" + if declared is None: + return DEFAULT_BEDROCK_REGION + if not declared.startswith(ENV_REFERENCE_PREFIX): + return declared + return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -24,7 +44,8 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if BEDROCK_ANTHROPIC_INFIX not in model: return None - return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + region: Final = bedrock_region(params.aws_region_name, model) + return None if region is None else f"bedrock/{region}" def route_bedrock( From 30c6241e3a5f534037dc57a6ea544d5ff9d8bdeb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:01:07 -0700 Subject: [PATCH 28/89] fix(e2e): a null error field is not an error Every OpenAI Responses body carries `error: null` at the top level, and the completeness check tested the key's presence rather than its value, so it rejected every single one. The cost was silent: nothing failed, the endpoint simply never cached, which is exactly the outcome the endpoint was added for. Found by driving the edge against the real providers rather than the synthetic fixtures, which carried no error key at all. Reading the value instead of the key is also more accurate for chat completions and messages, where a real error body carries a populated error object. --- .../test_provider_cache.py | 42 +++++++++++++++++-- tests/e2e/provider_cache.py | 7 +++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 0278bf8fd48..3dafbadf508 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -508,10 +508,13 @@ EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' ) -RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_SUCCESS: Final = ( + b'{"id":"resp_synthetic","object":"response","status":"completed","error":null,' + b'"incomplete_details":null,"output":[]}' +) RESPONSE_STREAM_SUCCESS: Final = ( - b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' - b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' + b'data: {"type":"response.created","response":{"id":"resp_synthetic","error":null}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"},"error":null}\n\n' ) @@ -586,6 +589,39 @@ class TestNonChatOpenAiEndpoints: assert call(url, MARKED).body == payload assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"id":"x","error":null,"choices":[{"message":{"content":"hi"},' + b'"finish_reason":"stop"}]}'), + ("/v1/messages", b'{"id":"msg_x","type":"message","role":"assistant","error":null,' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}'), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_a_null_error_field_is_not_an_error( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + """Every OpenAI Responses body carries `error: null`, and testing the key's + presence rather than its value rejected all of them. The cost was silent: + nothing failed, the endpoint simply never cached.""" + assert b'"error":null' in response + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"error":{"message":"rate limited","type":"rate_limit_error"}}'), + ("/v1/responses", b'{"object":"response","status":"completed","error":{"message":"bad"},"output":[]}'), + ]) + def test_a_populated_error_field_still_rejects( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,cacheable", [ ("/v1/chat/completions", True), ("/v1/messages", True), ("/v1/embeddings", True), ("/v1/responses", True), diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index a2b18a3a466..0dee33f33c6 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -167,7 +167,10 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]") except (UnicodeDecodeError, ValidationError): return False - if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): + if not values or any( + not isinstance(value, dict) or value.get("error") is not None or value.get("type") == "error" + for value in values + ): return False if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) @@ -187,7 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False - if not isinstance(value, dict) or "error" in value: + if not isinstance(value, dict) or value.get("error") is not None: return False path: Final = urlsplit(url).path if path == "/v1/messages": From c7246adc1dd571069a2cb3c77b3a40b720eb0da6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:29:02 -0700 Subject: [PATCH 29/89] fix(e2e): route only the Bedrock models the runner role can invoke The edge re-signs with the run pod's identity, whose IAM policy is an explicit per-model allowlist. Matching on the `anthropic.` infix instead routed every Anthropic-on-Bedrock model, so a model outside the policy came back 403 from Bedrock with no fallback, taking the whole claude_code Bedrock matrix red. An unlisted model now keeps its direct path and loses only caching. --- .../test_provider_cache.py | 4 ++++ tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/provider_cache_routing.py | 23 ++++++++++++------- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 3dafbadf508..ac1ccd5692b 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -809,6 +809,8 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), + ("bedrock/us.anthropic.claude-opus-4-7", "us-east-1"), + ("bedrock/converse/us.anthropic.claude-opus-4-7", "us-east-1"), ]) def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: """Almost every Bedrock deployment in the suite declares its region as @@ -839,6 +841,8 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-opus-4-5", aws_region_name="us-east-1"), + LiteLLMParamsBody(model="bedrock/converse/us.anthropic.claude-haiku-9-9", aws_region_name="us-east-1"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index c530574cda2..3a29c7b6fe4 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -22,7 +22,9 @@ Bedrock could not be mounted before because SigV4 signs the `Host` header, so a Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. -Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove + +Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index d2237bb49bc..97e05344423 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -9,8 +9,15 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) DEFAULT_BEDROCK_REGION: Final = "us-east-1" -BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." BEDROCK_CROSS_REGION_PREFIX: Final = "us." +BEDROCK_EDGE_MODELS: Final = frozenset( + { + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-5", + "us.anthropic.claude-opus-4-7", + } +) ENV_REFERENCE_PREFIX: Final = "os.environ/" @@ -33,16 +40,16 @@ def bedrock_region(declared: str | None, model: str) -> str | None: def bedrock_mount(params: LiteLLMParamsBody) -> str | None: - """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + """The edge mount a Bedrock deployment belongs to, or None. - Only the Anthropic models route. The edge validates converse and invoke - bodies by their Anthropic and Converse terminator fields, and the runner role - is allowed to invoke exactly those models, so Bedrock embeddings, image - generation, rerank and realtime keep their existing direct path rather than - reaching an edge that could neither sign nor validate for them.""" + The allowlist mirrors the runner role's IAM policy, which names its models + one by one. A model outside it would be re-signed with an identity that + cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps + its direct path and loses only caching. Adding a model is a policy edit in + litellm-ops and a line here.""" route: Final = params.model.partition("/")[2] model: Final = route.partition("/")[2] or route - if BEDROCK_ANTHROPIC_INFIX not in model: + if model not in BEDROCK_EDGE_MODELS: return None region: Final = bedrock_region(params.aws_region_name, model) return None if region is None else f"bedrock/{region}" From ebf34cd88006110c78a50648578e3dc7e0f115cd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:35:00 -0700 Subject: [PATCH 30/89] docs(e2e): say plainly that Bedrock streaming is not cached --- tests/e2e/PROVIDER_CACHE.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 3a29c7b6fe4..fe289e406aa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -2,7 +2,9 @@ `E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today ## Request identity From 7c2234be3a91a600c02f416e678b163dd53746ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 04:41:29 -0700 Subject: [PATCH 31/89] test(e2e): enforce the cross-region invariant on the Bedrock allowlist The allowlist rejects an unlisted model before the region resolver runs, so the two negative cases that used to cover the resolver were passing for the wrong reason and two mutations of it survived. Answering an env-referenced region with the default mount is only sound because every allowlisted model is a `us.` profile that fans out across the US regions, so assert that on the list itself and drop the per-call branch it made unreachable. --- .../test_provider_cache.py | 33 ++++++++++++++++++- tests/e2e/provider_cache_routing.py | 27 +++++++-------- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index ac1ccd5692b..c24ba8d6221 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -33,7 +33,13 @@ from provider_cache import ( successful_response, ) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store -from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model +from provider_cache_routing import ( + BEDROCK_CROSS_REGION_PREFIX, + BEDROCK_EDGE_MODELS, + LIVE_PROVIDER_REQUIRED, + bedrock_region, + route_cache_model, +) from fixture_mode import SESSION_TEST_KEY from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge from provider_edge_bedrock import bedrock_signer @@ -856,6 +862,31 @@ def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(par assert routed is params +@pytest.mark.parametrize("declared,expected", [ + (None, "us-east-1"), + ("us-west-2", "us-west-2"), + ("eu-west-1", "eu-west-1"), + ("os.environ/AWS_REGION", "us-east-1"), + ("os.environ/ANY_OTHER_NAME", "us-east-1"), +]) +def test_a_region_only_the_proxy_can_resolve_falls_back_to_the_default_mount( + declared: str | None, expected: str, +) -> None: + """A declared literal region is the one the deployment meant. A region the + proxy resolves from its own environment is one the run pod cannot see, and + the default mount answers it.""" + assert bedrock_region(declared) == expected + + +def test_every_model_on_the_edge_allowlist_is_a_cross_region_profile() -> None: + """Answering an env-referenced region with the default mount is only correct + for a profile that fans out across the US regions and is reachable from any + of them. A single-region model on this list would be sent to a region it may + not exist in, so the list is where that is caught.""" + assert BEDROCK_EDGE_MODELS + assert all(model.startswith(BEDROCK_CROSS_REGION_PREFIX) for model in BEDROCK_EDGE_MODELS) + + @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 97e05344423..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -21,22 +21,18 @@ BEDROCK_EDGE_MODELS: Final = frozenset( ENV_REFERENCE_PREFIX: Final = "os.environ/" -def bedrock_region(declared: str | None, model: str) -> str | None: - """The region whose edge mount a deployment belongs to, or None when the - harness cannot know it. +def bedrock_region(declared: str | None) -> str: + """The region whose edge mount a deployment belongs to. - Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy - resolves from its own environment. The run pod does not share that - environment, so the harness genuinely does not know the region. A `us.` - inference profile fans out across the US regions and is reachable from any - of them, so the default entry point is correct for those whatever the proxy - resolved; anything else keeps its direct path rather than being sent to a - region the model may not exist in.""" - if declared is None: + Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the + proxy can resolve from its own environment; the run pod does not share it. + Answering those with the default mount is correct because every model on the + edge allowlist is a `us.` inference profile, which fans out across the US + regions and is reachable from any of them. That invariant is enforced on the + allowlist itself rather than re-checked per call.""" + if declared is None or declared.startswith(ENV_REFERENCE_PREFIX): return DEFAULT_BEDROCK_REGION - if not declared.startswith(ENV_REFERENCE_PREFIX): - return declared - return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None + return declared def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -51,8 +47,7 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if model not in BEDROCK_EDGE_MODELS: return None - region: Final = bedrock_region(params.aws_region_name, model) - return None if region is None else f"bedrock/{region}" + return f"bedrock/{bedrock_region(params.aws_region_name)}" def route_bedrock( From bd1c2d6f07d7b9edd46bb11d95ad22cb7e029ea9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 05:09:04 -0700 Subject: [PATCH 32/89] fix(e2e): keep the tool-continuation echo-back test on the live path The key normalizes a unique marker so two builds match, which is the whole point, but it makes this test's identity collide with an earlier run's: it mints a fresh receipt, sends it through a tool result, and asserts the model echoes it back verbatim, so a stale recording matched and answered with the old receipt. Build 223 is where that surfaced, once the corpus was full enough for the first call to hit. A test that asserts a provider echoed this run's own unique value belongs on the live path. --- tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/conftest.py | 6 +++++- tests/e2e/llm_translation/test_messages_e2e.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index fe289e406aa..698e78a5aa1 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -44,7 +44,9 @@ The trusted runner receives: Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits -Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay +Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. + +One more class needs it, and it is the cost of normalizing the marker. A test that mints a fresh marker, sends it, and then asserts the provider's answer contains that exact value is asserting on the marker rather than using it as a salt. The key treats two such requests as the same identity, so a stale recording matches and answers with the marker from the run that recorded it. `TestOpenAIMessagesToolContinuation` is the one in the suite today: it sends a freshly minted receipt through a tool result and asserts the model echoes it back verbatim. If you add a test that asserts a provider echoed your own unique value, it belongs on the live path. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay ## Recorded response semantics diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 829c84910a9..430e16525d5 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -85,7 +85,11 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache") + config.addinivalue_line( + "markers", + "provider_live: requires actual provider timing, limits, state, or a response that echoes this" + " run's own unique value; bypass shared cache", + ) config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44c416a3e78..09ec48daa2f 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -372,6 +372,7 @@ def _request_tool( class TestOpenAIMessagesToolContinuation: + @pytest.mark.provider_live @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) def test_required_tool_arguments_and_correlated_result( self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool From 30a691ed55c59d8fd19c28e5356ba196bf5ae45a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 06:46:19 -0700 Subject: [PATCH 33/89] feat(e2e): cache Bedrock streaming responses The Claude Code compat cells drive the real CLI, which always streams, so converse-stream and invoke-with-response-stream were most of the suite's Bedrock traffic and all of it bypassed the edge. AWS frames those as binary vnd.amazon.eventstream rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. Two details drove the rule. A ConverseStream ends with metadata, not with messageStop, and metadata is what carries the token usage litellm prices the call from, so a stream cut between the two names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser: it yields the frames it did receive and silently discards a trailing partial one, so a stream cut one byte short parses clean. The body is checked against the frame lengths it declares to catch that. The invoke stream carries the ordinary Anthropic event grammar inside its chunk frames, so it shares the completeness rule with the SSE mounts. Validated against three real Bedrock eventstream captures, and the tests build their own frames rather than pasting a capture, with one test holding that framing to botocore's parser. --- .../test_provider_cache.py | 258 ++++++++++++++++-- tests/e2e/PROVIDER_CACHE.md | 4 +- tests/e2e/provider_cache.py | 141 +++++++++- 3 files changed, 370 insertions(+), 33 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index c24ba8d6221..4c131434ecd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -1,9 +1,13 @@ from __future__ import annotations +import base64 +import binascii +import json import os import shutil import socket import subprocess +import struct import threading import time import uuid @@ -17,9 +21,11 @@ from typing import Final from urllib.parse import urlsplit import pytest +from pydantic import JsonValue from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward from models import LiteLLMParamsBody, ModelMode from botocore.credentials import Credentials +from botocore.eventstream import EventStreamBuffer from provider_cache import ( SIGNATURE_HEADERS, CacheEdge, @@ -638,6 +644,53 @@ class TestNonChatOpenAiEndpoints: assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable +BEDROCK_STREAM_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" +CONVERSE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/converse-stream" +INVOKE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/invoke-with-response-stream" + + +def eventstream_frame(headers: Mapping[str, str], payload: bytes) -> bytes: + """AWS eventstream wire framing, the shape `vnd.amazon.eventstream` bodies + arrive in. Built here rather than pasted from a capture so a test can express + the stream it means; `test_the_frames_these_tests_build_are_real_aws_framing` + holds it to botocore's own parser.""" + encoded: Final = b"".join( + bytes([len(name)]) + name.encode() + b"\x07" + struct.pack(">H", len(value)) + value.encode() + for name, value in headers.items() + ) + prelude: Final = struct.pack(">II", 16 + len(encoded) + len(payload), len(encoded)) + framed: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + encoded + payload + return framed + struct.pack(">I", binascii.crc32(framed)) + + +def eventstream_event(event_type: str, payload: JsonValue, message_type: str = "event") -> bytes: + return eventstream_frame( + {":event-type": event_type, ":message-type": message_type, ":content-type": "application/json"}, + json.dumps(payload).encode(), + ) + + +def invoke_chunk(inner: JsonValue) -> bytes: + return eventstream_event("chunk", {"bytes": base64.b64encode(json.dumps(inner).encode()).decode("ascii")}) + + +CONVERSE_STREAM_OK: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("contentBlockStop", {"contentBlockIndex": 0}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + + eventstream_event("metadata", {"usage": {"inputTokens": 12, "outputTokens": 6, "totalTokens": 18}}) +) +INVOKE_STREAM_OK: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x", "role": "assistant"}}) + + invoke_chunk({"type": "content_block_start", "index": 0}) + + invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"text": "hi"}}) + + invoke_chunk({"type": "content_block_stop", "index": 0}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}) +) + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -738,32 +791,55 @@ class TestBedrockSigning: assert call(url, BEDROCK_BODY).body == response assert len(provider.hits) == 2 - @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) - def test_streaming_endpoints_go_live_every_time( - self, store: RedisResponseStore, provider: Provider, action: str, + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK), + ("invoke-with-response-stream", INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_finished_stream_is_served_from_the_cache_the_second_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, ) -> None: - """An eventstream's completeness cannot be proven without parsing its - frames, so these bypass rather than risk recording a truncated answer. - They are still signed: a bypass is a forward, not a passthrough.""" - provider.response = CONVERSE_SUCCESS - cache: Final = bedrock_cache_edge(store) - for _ in range(2): - with bedrock_edge(cache, provider, action) as url: - assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS - assert len(provider.hits) == 2 - assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:hits"] == 1 assert all( sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") for sent in provider.authorizations ), provider.authorizations - @pytest.mark.parametrize("action,cacheable", [ - ("converse", True), ("invoke", True), - ("converse-stream", False), ("invoke-with-response-stream", False), - ]) - def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK[:-1]), + ("invoke-with-response-stream", INVOKE_STREAM_OK[:-1]), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_the_connection_cut_short_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + """The whole risk of caching an eventstream is recording a half-finished + one, so a truncated body has to be rejected rather than stored.""" + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:rejected"] == 1 + assert f"mount:{BEDROCK_MOUNT}:hits" not in dict(replay.counters.counts) + + @pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream", "invoke-with-response-stream"]) + def test_every_anthropic_bedrock_action_is_cacheable(self, action: str) -> None: url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" - assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) + + @pytest.mark.parametrize("action", ["count-tokens", "invoke-async", "converse-stream-x"]) + def test_an_unknown_bedrock_action_is_not_cacheable(self, action: str) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert not cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) def test_a_region_mount_resolves_whole(self) -> None: resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) @@ -1021,3 +1097,147 @@ def test_duplicate_headers_bypass_cache_and_count_live_calls( assert len(provider.hits) == (2 if known_mount else 0) assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2 assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0) + + +class TestBedrockStreams: + def test_the_frames_these_tests_build_are_real_aws_framing(self) -> None: + buffer: Final = EventStreamBuffer() + buffer.add_data(CONVERSE_STREAM_OK) + assert [event.headers[":event-type"] for event in buffer] == [ + "messageStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata", + ] + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_finished_stream_is_recordable(self, url: str, body: bytes) -> None: + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, b"{}") + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + @pytest.mark.parametrize("keep", [1, -1, -4]) + def test_a_stream_the_connection_cut_short_is_not_recordable( + self, url: str, body: bytes, keep: int, + ) -> None: + """botocore yields the frames it did receive and silently drops a trailing + partial one, so a stream cut a single byte short parses clean and only the + byte accounting and the terminator rule catch it.""" + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body[:keep]) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_corrupted_frame_is_not_recordable(self, url: str, body: bytes) -> None: + flipped: Final = bytearray(body) + flipped[len(body) // 2] ^= 0xFF + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, bytes(flipped)) + + def test_a_converse_stream_that_lost_its_usage_is_not_recordable(self) -> None: + """ConverseStream names its stop reason a frame before it reports usage, + and litellm prices the call from that usage, so a stream cut between the + two would replay as a free call.""" + without_metadata: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + ) + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, without_metadata) + + def test_a_converse_stream_that_never_stopped_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_a_stream_that_failed_after_answering_200_is_not_recordable(self) -> None: + """Bedrock reports a fault that began after the headers went out as an + exception frame in place of the terminator it never got to send.""" + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("modelStreamErrorException", {"message": "boom"}, message_type="exception"), + ) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_cut_after_its_terminator_is_not_recordable(self, url: str, body: bytes) -> None: + """The terminator rules cannot see this one. Every frame the stream owes + has arrived and the partial frame after them is the one botocore drops + without a word, so only counting the bytes against the frame lengths + tells this from a stream that ended where it meant to.""" + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body + b"\x00\x00\x02") + + def test_a_converse_stream_whose_stop_frame_names_no_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_an_invoke_stream_carrying_a_frame_that_is_not_a_chunk_is_not_recordable(self) -> None: + """Every frame of an invoke stream is a `chunk` holding one base64 event. + A frame that is not one carries an event this rule cannot read, so the + stream can no longer be judged complete.""" + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_a_frame_claiming_no_length_is_rejected_rather_than_walked_forever(self) -> None: + """A frame length of zero never advances the cursor. Rejecting it is what + keeps a corrupt body from spinning the edge instead of answering.""" + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, b"\x00\x00\x00\x00" * 4) + + @pytest.mark.parametrize("url,terminator", [ + (INVOKE_STREAM_URL, invoke_chunk({"type": "message_stop"})), + (CONVERSE_STREAM_URL, eventstream_event("metadata", {"usage": {"totalTokens": 18}})), + ], ids=["invoke-stream", "converse-stream"]) + def test_a_delta_that_names_no_stop_reason_does_not_finish_a_stream( + self, url: str, terminator: bytes, + ) -> None: + """A `message_delta` arriving without its stop reason is the shape of a + turn the connection cut short partway through the delta itself.""" + head: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_delta", "delta": {}}) + ) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, head + terminator) + + def test_an_invoke_chunk_that_is_not_base64_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("chunk", {"bytes": "not base64 at all !!"}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_invoke_stream_missing_its_stop_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_empty_stream_is_not_recordable(self) -> None: + for url in (CONVERSE_STREAM_URL, INVOKE_STREAM_URL): + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, b"") + + def test_each_streaming_endpoint_is_held_to_its_own_grammar(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, INVOKE_STREAM_OK) + assert not successful_response(BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, CONVERSE_STREAM_OK) + + def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 698e78a5aa1..29d6e5ab4f2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -4,7 +4,9 @@ The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored -Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic + +Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way ## Request identity diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0dee33f33c6..9ae89861f63 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -13,6 +13,7 @@ from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit +from botocore.eventstream import EventStreamBuffer, ParserError from e2e_http import ( NetworkError, StreamChunk, @@ -36,6 +37,19 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +BEDROCK_CONVERSE_SUFFIX: Final = "/converse" +BEDROCK_INVOKE_SUFFIX: Final = "/invoke" +BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" +BEDROCK_INVOKE_STREAM_SUFFIX: Final = "/invoke-with-response-stream" +BEDROCK_SUFFIXES: Final = ( + BEDROCK_CONVERSE_SUFFIX, + BEDROCK_INVOKE_SUFFIX, + BEDROCK_CONVERSE_STREAM_SUFFIX, + BEDROCK_INVOKE_STREAM_SUFFIX, +) +EVENTSTREAM_PRELUDE_BYTES: Final = 4 +EVENT_TYPE_HEADER: Final = ":event-type" +EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -145,7 +159,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> return False path: Final = urlsplit(url).path if is_bedrock(mount): - return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) return path in OPENAI_JSON_PATHS @@ -176,16 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) - return ( - "[DONE]" not in events - and isinstance(values[0], dict) and values[0].get("type") == "message_start" - and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop" - and any( - isinstance(value, dict) and value.get("type") == "message_delta" - and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) - for value in values - ) - ) + return "[DONE]" not in events and complete_anthropic_stream(values) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: @@ -214,14 +219,19 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated or error body is missing the terminator field, which is what makes it safe to - record. The streaming variants never reach here: they are not cacheable.""" + record.""" + path: Final = urlsplit(url).path + if path.endswith(BEDROCK_CONVERSE_STREAM_SUFFIX): + return complete_converse_stream(body) + if path.endswith(BEDROCK_INVOKE_STREAM_SUFFIX): + return complete_invoke_stream(body) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False if not isinstance(value, dict) or "message" in value: return False - if urlsplit(url).path.endswith("/converse"): + if path.endswith(BEDROCK_CONVERSE_SUFFIX): return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) return ( value.get("type") == "message" @@ -230,6 +240,111 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def whole_eventstream_messages(body: bytes) -> bool: + """Whether the body is exactly a whole number of eventstream messages. + + A dropped connection is the failure this catches, and it has to be caught + here: botocore yields the messages it did receive and silently discards a + trailing partial one, so a stream cut a single byte short parses clean. Each + message declares its own total length in its first four bytes, so walking + those is enough to tell a complete body from a cut one.""" + offset = 0 # rebind-ok: a cursor walking the declared frame lengths + while offset + EVENTSTREAM_PRELUDE_BYTES <= len(body): + total: int = int.from_bytes(body[offset : offset + EVENTSTREAM_PRELUDE_BYTES], "big") + if total <= 0 or offset + total > len(body): + return False + offset += total + return offset == len(body) + + +def eventstream_events(body: bytes) -> tuple[tuple[str, JsonValue], ...] | None: + """The stream's (event type, decoded payload) pairs, or None if it is not a + complete, uncorrupted stream. + + botocore validates both CRCs and raises ``ParserError`` rather than decoding + corruption into something plausible. A failure that began after Bedrock had + already answered 200 arrives as an ``exception`` frame in place of the + terminator, so it is the terminator rules below that reject it and this does + not need to inspect ``:message-type`` as well.""" + if not body or not whole_eventstream_messages(body): + return None + buffer: Final = EventStreamBuffer() + buffer.add_data(body) + try: + return tuple( + (event_type(event.headers), JSON_VALUE.validate_json(event.payload)) + for event in buffer + ) + except (ParserError, ValidationError, ValueError): + return None + + +def event_type(headers: object) -> str: + """botocore's eventstream headers come back untyped, so the one header this + reads is validated into a string rather than trusted.""" + parsed: Final = EVENTSTREAM_HEADERS.validate_python(headers) + return parsed.get(EVENT_TYPE_HEADER, "") + + +def complete_converse_stream(body: bytes) -> bool: + """ConverseStream ends with ``metadata``, not with ``messageStop``. + + Requiring the metadata frame rather than the stop frame is deliberate: it + carries the token usage litellm prices the call from, so a stream cut between + the two still names a stop reason but would replay as a free call.""" + events: Final = eventstream_events(body) + if not events or events[-1][0] != "metadata": + return False + return any( + event_type == "messageStop" and isinstance(payload, dict) and isinstance(payload.get("stopReason"), str) + for event_type, payload in events + ) + + +def complete_invoke_stream(body: bytes) -> bool: + """InvokeModelWithResponseStream wraps the ordinary Anthropic event grammar + in ``chunk`` frames, one base64 payload each, so it is held to the same + terminator rule as the Anthropic SSE path. A frame Bedrock sends instead of a + chunk, an exception among them, carries no such payload and fails the rule + without the frame type needing to be read.""" + events: Final = eventstream_events(body) + if not events: + return False + values: Final = tuple(invoke_chunk_value(payload) for _, payload in events) + return all(value is not None for value in values) and complete_anthropic_stream(values) + + +def invoke_chunk_value(payload: JsonValue) -> JsonValue | None: + """The Anthropic event inside one ``chunk`` frame, or None for a frame that + carries no readable one.""" + if not isinstance(payload, dict) or not isinstance(encoded := payload.get("bytes"), str): + return None + try: + return JSON_VALUE.validate_json(base64.b64decode(encoded, validate=True)) + except (ValidationError, ValueError): + return None + + +def complete_anthropic_stream(values: tuple[JsonValue, ...]) -> bool: + """The Anthropic event grammar, shared by the SSE mounts and by Bedrock's + invoke stream, which carries the same events inside eventstream frames. A + ``message_delta`` naming a stop reason is what separates a finished turn from + one the connection cut short.""" + if not values: + return False + first: Final = values[0] + last: Final = values[-1] + return ( + isinstance(first, dict) and first.get("type") == "message_start" + and isinstance(last, dict) and last.get("type") == "message_stop" + and any( + isinstance(value, dict) and value.get("type") == "message_delta" + and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) + for value in values + ) + ) + + def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: """The Responses API streams typed events and ends with ``response.completed``. A run that failed, was cancelled, or ran out of tokens ends with a different From 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:38:44 -0700 Subject: [PATCH 34/89] feat(e2e): mount Gemini on the provider cache Gemini needs none of the machinery Bedrock needed. litellm composes {api_base}/models/{model}:{endpoint} from a custom api_base, so a plain path-prefixed mount reaches it, and the credential travels as a static x-goog-api-key header that no host rewrite invalidates. Nothing is re-signed and nothing leaves the cache key, so a recording still cannot cross credentials. A finished turn names a finishReason on every candidate and reports usageMetadata. The reason is read as a string rather than compared to STOP: MAX_TOKENS and the safety reasons end a turn just as finally, and rejecting them would send every one of them upstream forever. Streaming is the half worth care. Gemini repeats usageMetadata on every chunk and names a finishReason only on the last, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk carrying usage and no reason. The mount's upstream base carries the API version, so the path the rules see is /v1beta/models/..., not the one the proxy sent. The first version of this anchored the rule at the start of that path, which passed every test against a stub with no version prefix and would have cached nothing at all in a real run. Caught by replaying the rules over responses captured from live gemini-2.5-flash, which is also why the tests now mount their stub under the version prefix. Vertex stays unmounted and is a separate provider here: litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so Vertex needs a root-mounted edge on its own port. --- .../test_provider_cache.py | 154 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 16 +- tests/e2e/provider_cache.py | 46 ++++++ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 + 5 files changed, 215 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 4c131434ecd..520317a0678 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -861,7 +861,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -873,6 +873,8 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), + LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), + LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1241,3 +1243,153 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) + + +GEMINI_MODEL: Final = "gemini-2.5-flash" +GEMINI_API_VERSION: Final = "/v1beta" +GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" +GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" +GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} + + +def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: + candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} + return { + "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] + if candidates else [], + **({"usageMetadata": GEMINI_USAGE} if usage else {}), + "modelVersion": GEMINI_MODEL, + } + + +def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: + return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() + + +def gemini_stream(*finish_reasons: str | None) -> bytes: + return b"".join( + b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons + ) + + +@contextmanager +def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" + running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) + try: + yield running.edge.api_base("gemini") + path + finally: + running.shutdown() + + +class TestGemini: + """Gemini reaches the edge by path prefix alone: litellm composes + `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, + so nothing has to be re-signed and nothing leaves the cache key. The response + grammar is its own though, and the streaming one is the interesting half: every + chunk repeats `usageMetadata`, so only `finishReason` on the last chunk + separates a finished turn from a dropped connection.""" + + @pytest.mark.parametrize("path,response", [ + (GEMINI_GENERATE_PATH, gemini_unary()), + (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), + ], ids=["generate", "stream"]) + def test_a_finished_turn_replays_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.stream = path == GEMINI_STREAM_PATH + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) + def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( + self, store: RedisResponseStore, provider: Provider, reason: str, + ) -> None: + """Reading `finishReason` as a string rather than comparing it to STOP is + deliberate. A turn cut off by the token limit or a safety filter is over, + and rejecting those would send every one of them upstream forever.""" + provider.response = gemini_unary(reason) + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("response", [ + gemini_unary(None), + gemini_unary("STOP", usage=False), + gemini_unary("STOP", candidates=False), + b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', + ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) + def test_an_unfinished_or_failed_turn_never_enters_the_cache( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("response", [ + gemini_stream(None, None), + gemini_stream("STOP", None), + gemini_stream(), + ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) + def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.stream = True + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A request for more than one candidate is answered by more than one, and + the turn is over only when every one of them names a reason. Holding the + whole list to that rule rather than its first entry is what keeps a + half-finished answer from being stored and replayed as a finished one.""" + finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] + unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] + provider.response = json.dumps( + {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} + ).encode() + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + (GEMINI_GENERATE_PATH, True), + (GEMINI_STREAM_PATH, True), + (f"/models/{GEMINI_MODEL}:countTokens", False), + (f"/models/{GEMINI_MODEL}:embedContent", False), + ("/v1/chat/completions", False), + (f"/files/{GEMINI_MODEL}:generateContent", False), + ]) + @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) + def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: + """The mount's upstream base carries the API version, so the path the cache + sees is the upstream one and starts `/v1beta`. A rule anchored at the start + of the path would pass every test against a stub with no version prefix and + then cache nothing at all in a real run.""" + assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable + + def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: + """The shapes above are hand-built so a test can express the turn it means. + This holds them to the fields a live `generativelanguage.googleapis.com` + answer carries, captured 2026-09-16 against gemini-2.5-flash.""" + captured: Final = json.loads( + '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' + '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' + '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' + ) + built: Final = json.loads(gemini_unary()) + assert captured.keys() >= built.keys() + assert captured["candidates"][0].keys() >= built["candidates"][0].keys() + assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 29d6e5ab4f2..e2f8030074d 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,13 +1,23 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way +## Gemini + +Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials + +The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run + +A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason + +Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments + ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -30,7 +40,7 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm +Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 9ae89861f63..7bba93321b5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,6 +37,10 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +GEMINI_MOUNT: Final = "gemini" +GEMINI_MODELS_SEGMENT: Final = "/models" +GEMINI_GENERATE_SUFFIX: Final = ":generateContent" +GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -154,12 +158,21 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX +def is_gemini(mount: str) -> bool: + return mount == GEMINI_MOUNT + + def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) + if is_gemini(mount): + collection, _, resource = path.rpartition("/") + return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( + (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) + ) return path in OPENAI_JSON_PATHS @@ -186,6 +199,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False + if is_gemini(mount): + return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -197,6 +212,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False + if is_gemini(mount): + return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -215,6 +232,35 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) +def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: + """A finished Gemini turn names a ``finishReason`` on every candidate and + reports the usage litellm prices the call from. ``finishReason`` is read as a + string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety + reasons end a turn just as finally, and a cache that rejected them would send + every one of them upstream forever.""" + candidates: Final = value.get("candidates") + return ( + isinstance(value.get("usageMetadata"), dict) + and isinstance(candidates, list) + and bool(candidates) + and all( + isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) + for candidate in candidates + ) + ) + + +def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: + """Gemini repeats ``usageMetadata`` on every chunk but names a + ``finishReason`` only on the last one, so the terminator is the final event + rather than any event. A stream the connection cut short ends on a chunk that + carries usage and no reason, which is exactly what this rejects.""" + if not values: + return False + last: Final = values[-1] + return isinstance(last, dict) and complete_gemini_candidates(last) + + def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index f9775a2b152..c4b02beac2d 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,6 +19,7 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" +EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -81,7 +82,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in {"openai", "anthropic"}: + if provider not in EDGE_PROVIDERS: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 2606b26fe99..219df55233a 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,6 +104,7 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 7006da9cde950981379214f3d2dd0f645d68995e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:47:31 -0700 Subject: [PATCH 35/89] feat(e2e): say why a response was not recorded Build 226 routed Bedrock streaming for the first time and rejected 62 of 220 misses on that mount, and the counters could not say why. A flat rejected count covers three unrelated things with opposite fixes: the consumer walking away mid-capture, a body that arrived whole and failed its endpoint's rule, and a provider that could not be reached. Each now also counts its own reason. A consumer that walks away was counting nothing at all. Abandoning the capture generator raises GeneratorExit at its yield, so neither branch of the old accounting ran and the miss simply vanished from the report, which is also why misses could exceed writes plus rejected with nothing to explain the gap. The decision moves into settle() so the generator's finally owns the accounting and an abandoned capture is counted like any other rejection. --- .../test_provider_cache.py | 39 ++++++++++++++ tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 53 +++++++++++++------ 3 files changed, 78 insertions(+), 16 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 520317a0678..1271b20438a 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -516,6 +516,45 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( + store: RedisResponseStore, provider: Provider, +) -> None: + """One `rejected` count cannot tell a connection that dropped from a body the + provider finished sending and the rules turned down, and those have opposite + fixes: the first is the client going away mid-capture, the second is a grammar + the cache does not accept. A mount whose rejections are mostly one or the other + is a different problem, so the report has to be able to say which.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cut_short: Final = cache_edge(store) + provider.stream = True + provider.truncated = True + provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' + running: Final = start_provider_edge(cut_short, mounts={"openai": upstream}) + try: + forward("POST", running.edge.api_base("openai") + "/v1/chat/completions", + headers=HEADERS, body=MARKED, timeout=5) + finally: + running.shutdown() + + unfinished: Final = cache_edge(store) + provider.stream = False + provider.truncated = False + provider.response = b'{"choices":[{"index":0,"message":{"content":"hi"}}]}' + second: Final = start_provider_edge(unfinished, mounts={"openai": upstream}) + try: + call(second.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + second.shutdown() + + cut: Final = dict(cut_short.counters.counts) + turned_down: Final = dict(unfinished.counters.counts) + assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + assert cut["mount:openai:rejected_cut_short"] == 1 + assert "mount:openai:rejected_incomplete" not in cut + assert turned_down["mount:openai:rejected_incomplete"] == 1 + assert "mount:openai:rejected_cut_short" not in turned_down + + EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e2f8030074d..894d9be4efa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 7bba93321b5..f528d08b720 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -52,6 +52,9 @@ BEDROCK_SUFFIXES: Final = ( BEDROCK_INVOKE_STREAM_SUFFIX, ) EVENTSTREAM_PRELUDE_BYTES: Final = 4 +CUT_SHORT: Final = "cut_short" +INCOMPLETE: Final = "incomplete" +UNREACHABLE: Final = "unreachable" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -551,7 +554,7 @@ class CacheEdge: ) prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return prepared identity: Final = request_identity( self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, @@ -575,7 +578,7 @@ class CacheEdge: return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return head return StreamHead( head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), @@ -585,25 +588,45 @@ class CacheEdge: self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() + reason = CUT_SHORT # rebind-ok: a consumer that walks away never reaches the settle call below try: with closing(head.steps): yield StreamChunk(b"") for step in head.steps: yield step capture.observe(step) - chunks: Final = capture.chunks() if capture.eligible else () - headers: Final = { - name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS - } - if not capture.eligible or not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): - self.count(mount, "rejected") - return - response: Final = CachedResponse( - request_key=key, status_code=head.status_code, headers=headers, - chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), - ) - published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) - self.count(mount, "writes" if published else "write_failures") + reason = self.settle(mount, key, lease, url, head, capture) finally: + self.reject(mount, reason) self.store.release(key, lease) capture.buffer.close() + + def settle( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, capture: ResponseCapture, + ) -> str | None: + """None once the response is stored, otherwise the reason it was not.""" + if not capture.eligible: + return CUT_SHORT + headers: Final = { + name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS + } + chunks: Final = capture.chunks() + if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + return INCOMPLETE + response: Final = CachedResponse( + request_key=key, status_code=head.status_code, headers=headers, + chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), + ) + published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) + self.count(mount, "writes" if published else "write_failures") + return None + + def reject(self, mount: str, reason: str | None) -> None: + """A flat rejection count cannot separate a connection that went away from + a body the provider finished sending and the rules turned down, and the two + have opposite fixes. A mount whose rejections are nearly all one or the + other is a different problem, so the report has to be able to say which.""" + if reason is None: + return + self.count(mount, "rejected") + self.count(mount, f"rejected_{reason}") From fdb8e3533be56df5de394bdb155870008c50bf12 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:58:48 -0700 Subject: [PATCH 36/89] fix(mcp): validate credentials in existing request paths --- .../mcp_server/mcp_server_manager.py | 42 ++++----- .../mcp_server/openapi_to_mcp_generator.py | 40 +++++++-- .../outbound_credentials/adapter.py | 67 ++++++++++++++- .../proxy/_experimental/mcp_server/server.py | 1 - .../_experimental/mcp_server/upstream.py | 85 ------------------- .../proxy/_experimental/mcp_server/utils.py | 16 ---- .../mcp_server/test_mcp_hook_extra_headers.py | 1 - .../mcp_server/test_mcp_server_manager.py | 76 ++++++++++++++--- .../test_openapi_to_mcp_generator.py | 84 ++++++++++++++++++ 9 files changed, 265 insertions(+), 147 deletions(-) delete mode 100644 litellm/proxy/_experimental/mcp_server/upstream.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0254f79cbcc..6881956595c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( UpstreamCredentialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + prepare_mcp_client, raise_public, raise_token_exchange_challenge, raise_user_oauth_challenge, @@ -132,7 +133,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) -from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client, validate_openapi_credentials from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -2805,6 +2805,8 @@ class MCPServerManager: headers=headers, server_label=server.name or server.server_name or server.alias or server.server_id, relays_upstream_auth=server.is_client_forwarded_token, + auth_type=server.auth_type, + upstream_token_header=server.upstream_token_header, ) tool_func.__name__ = prefixed_tool_name tool_func.__doc__ = description @@ -4230,19 +4232,16 @@ class MCPServerManager: ) record_auth_resolution(server.server_id, AuthResolution.not_applicable) - return await prepare_mcp_client( - resolved_server, - MCPClient( - server_url="", # Not used for stdio - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - stdio_config=stdio_config, - extra_headers=extra_headers, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, - ), + return MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + stdio_config=stdio_config, + extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) else: # For HTTP/SSE transports @@ -6200,7 +6199,6 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, forwarded_headers: dict[str, str] | None, - caller_authorization: str | None = None, ) -> tuple[dict[str, str] | None, dict[str, str] | None]: """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. @@ -6224,12 +6222,9 @@ class MCPServerManager: """ spec: Final = to_server_spec(mcp_server) if spec is None: - stored_headers = ( - None - if oauth2_headers - else await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) - ) - validate_openapi_credentials(mcp_server, stored_headers, forwarded_headers, caller_authorization) + if oauth2_headers: + return None, forwarded_headers + stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) return stored_headers, forwarded_headers subject_token: str | None = None @@ -6248,9 +6243,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=forwarded_headers, ) - resolved_headers: Final = await _materialize_auth_headers(resolved_auth) - validate_openapi_credentials(mcp_server, resolved_headers, forwarded_headers, caller_authorization) - return resolved_headers, forwarded_headers + return await _materialize_auth_headers(resolved_auth), forwarded_headers async def _gather_openapi_tool_tasks( self, @@ -6376,7 +6369,6 @@ class MCPServerManager: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) async def _call_openapi_via_handler(): diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 66712e97a34..477d86ab436 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -20,7 +20,6 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, MCPUpstreamAuthError, ) -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to @@ -55,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) -from litellm.types.mcp import credential_redirect_hook, custom_credential_slot +from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot class _OpenAPIJSONSchema(TypedDict, total=False): @@ -416,9 +415,26 @@ def _merge_openapi_tool_request_headers( Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ - return merge_openapi_headers( - static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get() - ) + request_extra: Final = _request_extra_headers.get() or {} + static: Final = static_headers or {} + + static_lower_names: Final = {k.lower() for k in static} + effective_headers: dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} + effective_headers.update(static) + + override_auth: Final = _request_auth_header.get() + if override_auth: + for existing in [k for k in effective_headers if k.lower() == "authorization"]: + del effective_headers[existing] + effective_headers["Authorization"] = override_auth + + resolved_auth_headers: Final = _request_resolved_auth_headers.get() or {} + for name, value in resolved_auth_headers.items(): + for existing in [k for k in effective_headers if k.lower() == name.lower()]: + del effective_headers[existing] + effective_headers[name] = value + + return effective_headers def _raise_for_upstream_failure( @@ -455,6 +471,8 @@ def create_tool_function( headers: dict[str, str] | None = None, server_label: str | None = None, relays_upstream_auth: bool = False, + auth_type: MCPAuthType = None, + upstream_token_header: str | None = None, ): """Create a tool function for an OpenAPI operation. @@ -487,6 +505,18 @@ def create_tool_function( by using **kwargs instead of named parameters. """ effective_headers: Final = _merge_openapi_tool_request_headers(headers) + if auth_type is not None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_public, + validate_static_credential, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok + + match validate_static_credential(auth_type, effective_headers, upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + pass # Build URL from base_url and path url = base_url + path diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index d25946d81d0..5358878a248 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -13,15 +13,17 @@ from __future__ import annotations import base64 import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException from pydantic import SecretStr from typing_extensions import assert_never -from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials +from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, ApiKeyConfig, @@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( Subject, TokenExchangeConfig, ) -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth +from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth @@ -385,3 +387,64 @@ def raise_token_exchange_challenge( detail="Unauthorized", headers={"WWW-Authenticate": www_authenticate}, ) + + +_STATIC_MODES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) +) + + +def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: + if not value: + return False + if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + return True + if value.lower() in ("bearer", "basic", "token", "apikey"): + return False + if auth_type in (MCPAuth.bearer_token, MCPAuth.token): + scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" + credential: Final = strip_auth_scheme(value, scheme).strip() + return bool(credential) and credential.lower() != scheme.lower() + if auth_type == MCPAuth.basic: + parts: Final = value.split(None, 1) + if len(parts) != 2 or parts[0].lower() != "basic": + return False + try: + decoded: Final = base64.b64decode(parts[1], validate=True).strip() + return b":" in decoded + except ValueError: + return False + return True + + +def validate_static_credential( + auth_type: MCPAuthType, + headers: Mapping[str, str], + upstream_token_header: str | None = None, +) -> Result[None, CredError]: + if auth_type not in _STATIC_MODES: + return Ok(None) + default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization" + slots: Final = frozenset( + name.lower() + for name in ( + upstream_token_header or default_slot, + default_slot, + "Authorization", + ) + ) + values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if any(_usable_credential_value(auth_type, name, value) for name, value in values): + return Ok(None) + return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential")) + + +async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: + if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: + return client + request: Final = await client.prepare_request_auth() + match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + return client diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a3aaada41f7..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3141,7 +3141,6 @@ if MCP_AVAILABLE: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) _auth_token: Final = _request_auth_header.set(auth_header_value) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py deleted file mode 100644 index 66840db21ce..00000000000 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -import base64 -from collections.abc import Mapping -from typing import Final - -from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme -from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public -from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result -from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers -from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport -from litellm.types.mcp_server.mcp_server_manager import MCPServer - -_STATIC_MODES: Final = frozenset( - (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) -) - - -def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: - if not value: - return False - if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): - return True - if value.lower() in ("bearer", "basic", "token", "apikey"): - return False - if auth_type in (MCPAuth.bearer_token, MCPAuth.token): - scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" - credential: Final = strip_auth_scheme(value, scheme).strip() - return bool(credential) and credential.lower() != scheme.lower() - if auth_type == MCPAuth.basic: - parts: Final = value.split(None, 1) - if len(parts) != 2 or parts[0].lower() != "basic": - return False - try: - decoded: Final = base64.b64decode(parts[1], validate=True).strip() - return b":" in decoded - except ValueError: - return False - return True - - -def validate_static_credential(server: MCPServer, headers: Mapping[str, str]) -> Result[None, CredError]: - if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: - return Ok(None) - default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" - slots: Final = frozenset( - name.lower() - for name in ( - server.upstream_token_header or default_slot, - default_slot, - "Authorization", - ) - ) - values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) - if any(_usable_credential_value(server.auth_type, name, value) for name, value in values): - return Ok(None) - return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) - - -async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: - if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: - return client - request: Final = await client.prepare_request_auth() - match validate_static_credential(server, request.headers): - case Error(error): - raise_public(error) - case Ok(): - return client - - -def validate_openapi_credentials( - server: MCPServer, - resolved_headers: Mapping[str, str] | None, - forwarded_headers: Mapping[str, str] | None, - caller_authorization: str | None, -) -> None: - headers: Final = merge_openapi_headers( - server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers - ) - match validate_static_credential(server, headers): - case Error(error): - raise_public(error) - case Ok(): - return diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index bea74d36b34..fb3eb06fd15 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -756,22 +756,6 @@ def build_env_var_setup_url(server_id: str) -> str: return f"{base}{path}" if base else path -def merge_openapi_headers( - static_headers: Mapping[str, str], - extra_headers: Mapping[str, str] | None, - caller_authorization: str | None, - resolved_headers: Mapping[str, str] | None, -) -> dict[str, str]: - sources: Final = ( - extra_headers or {}, - static_headers, - {"Authorization": caller_authorization} if caller_authorization else {}, - resolved_headers or {}, - ) - entries: Final = {name.lower(): (name, value) for source in sources for name, value in source.items()} - return dict(entries.values()) - - def merge_mcp_headers( *, extra_headers: Mapping[str, str] | None = None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 5e3a26fb4ac..28faf375ab8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1391,7 +1391,6 @@ class TestOpenApiResolvedUpstreamAuth: mcp_auth_header="user-byok-key", user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), forwarded_headers=None, - caller_authorization="ApiKey user-byok-key", ) assert resolved is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 54add273c24..4c4c45162ca 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,11 +5,13 @@ import logging import os import sys from datetime import datetime +from pathlib import Path from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from respx import MockRouter from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -5127,7 +5129,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers captured["server_label"] = server_label @@ -5212,7 +5215,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers @@ -13471,6 +13475,41 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( class TestProtectedCredentialPreparation: + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,credential", [ + (MCPAuth.bearer_token, None), + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.api_key, None), + (MCPAuth.basic, "Basic"), + ]) + @pytest.mark.parametrize("dispatch", ["managed", "local"]) + async def test_openapi_dispatch_rejects_unusable_effective_credentials( + self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, credential: str | None, dispatch: str, + ) -> None: + from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool + from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix + + spec_path: Final = tmp_path / "openapi.json" + spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}})) + server: Final = MCPServer( + server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential, + ) + manager: Final = MCPServerManager() + await manager._register_openapi_tools(str(spec_path), server, server.url) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="unexpected success") + result: Final = ( + await manager._call_openapi_tool_handler(server, "echo", {}) + if dispatch == "managed" + else await _handle_local_mcp_tool(add_server_prefix_to_name("echo", get_server_prefix(server)), {}) + ) + assert result.isError is True + assert "requires a usable upstream credential" in result.content[0].text + assert destination.call_count == 0 + @pytest.mark.asyncio @pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse]) @pytest.mark.parametrize("client_secret", [None, ""]) @@ -13523,7 +13562,7 @@ class TestProtectedCredentialPreparation: assert client._get_auth_headers() == headers @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token]) + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange]) async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: server = MCPServer( server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", @@ -13594,22 +13633,35 @@ class TestProtectedCredentialPreparation: ({"X-API-Key": "static"}, {"Authorization": ""}, None), ]) async def test_openapi_static_credentials_remain_supported( - self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None ) -> None: - server = MCPServer(server_id="openapi-static", name="openapi-static", url="https://upstream.example", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static) - resolved, retained = await MCPServerManager().resolve_openapi_upstream_auth( - mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, - user_api_key_auth=None, forwarded_headers=forwarded, caller_authorization=caller, + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_auth_header, _request_extra_headers, create_tool_function, ) - assert resolved is None - assert retained == forwarded + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + try: + assert await tool() == "authenticated" + sent: Final = destination.calls.last.request.headers + assert sent.get("x-api-key") == static.get("X-API-Key", (forwarded or {}).get("X-API-Key")) + if caller: + assert sent["authorization"] == caller + assert destination.call_count == 1 + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) @pytest.mark.asyncio async def test_static_resolution_cancellation_closes_flow(self) -> None: from collections.abc import AsyncGenerator from litellm.experimental_mcp_client.client import MCPClient - from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import prepare_mcp_client class CancelledAuth(httpx.Auth): closed = False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 66c5627bc94..979199d0dc9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -10,9 +10,14 @@ This test suite ensures that: """ from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, patch import pytest +from fastapi import HTTPException +from respx import MockRouter + +from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, @@ -35,6 +40,85 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ + ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"), + ({"Authorization": "Bearer configured"}, None, "Bearer", None, None), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None), + ({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"), +]) +async def test_static_auth_validates_headers_after_existing_precedence( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None, + resolved: dict[str, str] | None, expected: str | None, +) -> None: + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + resolved_token: Final = _request_resolved_auth_headers.set(resolved) + try: + if expected is None: + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await tool() + assert exc.value.status_code == 500 + assert destination.call_count == 0 + else: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["authorization"] == expected + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("credential", ["custom-key", ""]) +async def test_static_auth_uses_configured_custom_header( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential}, + auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + if credential: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["x-custom"] == credential + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential"): + await tool() + assert destination.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type,resolved", [ + (MCPAuth.none, None), + (MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}), +]) +async def test_static_validation_preserves_no_auth_and_resolved_oauth( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, resolved: dict[str, str] | None, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="echo") + token: Final = _request_resolved_auth_headers.set(resolved) + try: + assert await tool() == "echo" + assert destination.call_count == 1 + assert destination.calls.last.request.headers.get("authorization") == (resolved or {}).get("Authorization") + finally: + _request_resolved_auth_headers.reset(token) + + def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock: """Utility to create a mocked async httpx client for the given method. From c447c3312db1b99e058062f8cd61f2904fa1e9ef Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:13:04 -0700 Subject: [PATCH 37/89] feat(e2e): separate a provider error from a body that failed its rule Build 226's 62 Bedrock rejections are the question this is trying to answer, and "incomplete" would have covered both candidate causes at once. Replaying the completeness rules over eight streams captured from live Bedrock, covering tool use, extended thinking and a max-tokens stop on both streaming endpoints, accepts every one of them, so a rule that is too strict is the less likely half. A provider that answered 429 or 5xx and was retried out of sight is the other, and it now counts as rejected_error_status rather than being folded in with a grammar failure. --- .../code_coverage_tests/test_provider_cache.py | 18 +++++++++++++++--- tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 3 +++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 1271b20438a..9668de35a02 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -546,13 +546,25 @@ def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( finally: second.shutdown() + refused: Final = cache_edge(store) + provider.status = 429 + provider.response = b'{"message":"Too many requests"}' + third: Final = start_provider_edge(refused, mounts={"openai": upstream}) + try: + call(third.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + third.shutdown() + cut: Final = dict(cut_short.counters.counts) turned_down: Final = dict(unfinished.counters.counts) - assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + errored: Final = dict(refused.counters.counts) + assert cut["mount:openai:rejected"] == turned_down["mount:openai:rejected"] == errored["mount:openai:rejected"] == 1 assert cut["mount:openai:rejected_cut_short"] == 1 - assert "mount:openai:rejected_incomplete" not in cut assert turned_down["mount:openai:rejected_incomplete"] == 1 - assert "mount:openai:rejected_cut_short" not in turned_down + assert errored["mount:openai:rejected_error_status"] == 1 + assert not {"mount:openai:rejected_incomplete", "mount:openai:rejected_error_status"} & set(cut) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_error_status"} & set(turned_down) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_incomplete"} & set(errored) EMBEDDING_SUCCESS: Final = ( diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 894d9be4efa..ed1e7d517b0 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index f528d08b720..399a0379889 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -55,6 +55,7 @@ EVENTSTREAM_PRELUDE_BYTES: Final = 4 CUT_SHORT: Final = "cut_short" INCOMPLETE: Final = "incomplete" UNREACHABLE: Final = "unreachable" +ERROR_STATUS: Final = "error_status" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -610,6 +611,8 @@ class CacheEdge: headers: Final = { name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS } + if not 200 <= head.status_code < 300: + return ERROR_STATUS chunks: Final = capture.chunks() if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): return INCOMPLETE From 1972a30defcea23d170ba431af99f2e82e652b24 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:37:25 -0700 Subject: [PATCH 38/89] revert(e2e): unmount Gemini, its api_base means two things Build 227 mounted Gemini and turned TestGeminiFiles::test_gemini_file_upload red. litellm's two Gemini endpoints disagree about what api_base means. Chat composes {api_base}/models/{model}:{endpoint} and defaults api_base to https://generativelanguage.googleapis.com/v1beta, so the version lives inside it. File upload composes {api_base}/upload/v1beta/files and defaults to the host root, so the version lives outside it. A single api_base cannot satisfy both, and a registration carries no signal about which endpoint the deployment will be used for, so the edge cannot route one and not the other. Backing it out rather than working around it. The cache must never turn a passing test red, which is the same rule the Bedrock model allowlist follows, and Gemini was 7 of roughly 1030 edge calls in that build. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits this too, so the fix belongs in litellm; mounting Gemini is one line once it lands. This reverts commit 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321. --- .../test_provider_cache.py | 154 +----------------- tests/e2e/PROVIDER_CACHE.md | 18 +- tests/e2e/provider_cache.py | 46 ------ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 - 5 files changed, 7 insertions(+), 215 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 9668de35a02..e57d33a0406 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -912,7 +912,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -924,8 +924,6 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), - LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), - LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1294,153 +1292,3 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) - - -GEMINI_MODEL: Final = "gemini-2.5-flash" -GEMINI_API_VERSION: Final = "/v1beta" -GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" -GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" -GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} - - -def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: - candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} - return { - "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] - if candidates else [], - **({"usageMetadata": GEMINI_USAGE} if usage else {}), - "modelVersion": GEMINI_MODEL, - } - - -def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: - return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() - - -def gemini_stream(*finish_reasons: str | None) -> bytes: - return b"".join( - b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons - ) - - -@contextmanager -def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: - upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" - running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) - try: - yield running.edge.api_base("gemini") + path - finally: - running.shutdown() - - -class TestGemini: - """Gemini reaches the edge by path prefix alone: litellm composes - `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, - so nothing has to be re-signed and nothing leaves the cache key. The response - grammar is its own though, and the streaming one is the interesting half: every - chunk repeats `usageMetadata`, so only `finishReason` on the last chunk - separates a finished turn from a dropped connection.""" - - @pytest.mark.parametrize("path,response", [ - (GEMINI_GENERATE_PATH, gemini_unary()), - (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), - ], ids=["generate", "stream"]) - def test_a_finished_turn_replays_on_the_next_run( - self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, - ) -> None: - provider.stream = path == GEMINI_STREAM_PATH - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, path) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) - def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( - self, store: RedisResponseStore, provider: Provider, reason: str, - ) -> None: - """Reading `finishReason` as a string rather than comparing it to STOP is - deliberate. A turn cut off by the token limit or a safety filter is over, - and rejecting those would send every one of them upstream forever.""" - provider.response = gemini_unary(reason) - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("response", [ - gemini_unary(None), - gemini_unary("STOP", usage=False), - gemini_unary("STOP", candidates=False), - b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', - ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) - def test_an_unfinished_or_failed_turn_never_enters_the_cache( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("response", [ - gemini_stream(None, None), - gemini_stream("STOP", None), - gemini_stream(), - ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) - def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.stream = True - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( - self, store: RedisResponseStore, provider: Provider, - ) -> None: - """A request for more than one candidate is answered by more than one, and - the turn is over only when every one of them names a reason. Holding the - whole list to that rule rather than its first entry is what keeps a - half-finished answer from being stored and replayed as a finished one.""" - finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] - unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] - provider.response = json.dumps( - {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} - ).encode() - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("path,cacheable", [ - (GEMINI_GENERATE_PATH, True), - (GEMINI_STREAM_PATH, True), - (f"/models/{GEMINI_MODEL}:countTokens", False), - (f"/models/{GEMINI_MODEL}:embedContent", False), - ("/v1/chat/completions", False), - (f"/files/{GEMINI_MODEL}:generateContent", False), - ]) - @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) - def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: - """The mount's upstream base carries the API version, so the path the cache - sees is the upstream one and starts `/v1beta`. A rule anchored at the start - of the path would pass every test against a stub with no version prefix and - then cache nothing at all in a real run.""" - assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable - - def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: - """The shapes above are hand-built so a test can express the turn it means. - This holds them to the fields a live `generativelanguage.googleapis.com` - answer carries, captured 2026-09-16 against gemini-2.5-flash.""" - captured: Final = json.loads( - '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' - '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' - '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' - ) - built: Final = json.loads(gemini_unary()) - assert captured.keys() >= built.keys() - assert captured["candidates"][0].keys() >= built["candidates"][0].keys() - assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index ed1e7d517b0..aca81f26c5a 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,23 +1,13 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way -## Gemini - -Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials - -The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run - -A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason - -Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments - ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -40,7 +30,9 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section +Vertex and Gemini are not mounted, for different reasons. litellm grafts the default Vertex path onto an `api_base` only when that `api_base` has no path of its own, so a path-prefixed Vertex mount instead becomes `{api_base}:{endpoint}`, dropping project, location and model. Vertex needs a root-mounted edge on its own port, or a change in litellm + +Gemini reaches a path-prefixed mount perfectly well and was mounted for one build, then backed out, because litellm's two Gemini endpoints disagree about what `api_base` means. Chat composes `{api_base}/models/{model}:{endpoint}` and defaults `api_base` to `https://generativelanguage.googleapis.com/v1beta`, so the version has to be inside it. File upload composes `{api_base}/upload/v1beta/files` and defaults to the host root, so the version has to be outside it. One `api_base` cannot satisfy both, and a deployment gives no signal at registration time about which it will be used for, so mounting Gemini turned `TestGeminiFiles::test_gemini_file_upload` red in build 227. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits the same thing; it is a litellm bug rather than a cache limitation, and mounting Gemini is one line once it is fixed Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 399a0379889..444972ffce5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,10 +37,6 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" -GEMINI_MOUNT: Final = "gemini" -GEMINI_MODELS_SEGMENT: Final = "/models" -GEMINI_GENERATE_SUFFIX: Final = ":generateContent" -GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -162,21 +158,12 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX -def is_gemini(mount: str) -> bool: - return mount == GEMINI_MOUNT - - def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) - if is_gemini(mount): - collection, _, resource = path.rpartition("/") - return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( - (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) - ) return path in OPENAI_JSON_PATHS @@ -203,8 +190,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False - if is_gemini(mount): - return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -216,8 +201,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False - if is_gemini(mount): - return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -236,35 +219,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) -def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: - """A finished Gemini turn names a ``finishReason`` on every candidate and - reports the usage litellm prices the call from. ``finishReason`` is read as a - string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety - reasons end a turn just as finally, and a cache that rejected them would send - every one of them upstream forever.""" - candidates: Final = value.get("candidates") - return ( - isinstance(value.get("usageMetadata"), dict) - and isinstance(candidates, list) - and bool(candidates) - and all( - isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) - for candidate in candidates - ) - ) - - -def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: - """Gemini repeats ``usageMetadata`` on every chunk but names a - ``finishReason`` only on the last one, so the terminator is the final event - rather than any event. A stream the connection cut short ends on a chunk that - carries usage and no reason, which is exactly what this rejects.""" - if not values: - return False - last: Final = values[-1] - return isinstance(last, dict) and complete_gemini_candidates(last) - - def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index c4b02beac2d..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,7 +19,6 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" -EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -82,7 +81,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in EDGE_PROVIDERS: + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 219df55233a..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,7 +104,6 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", - "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 6c517bfc49eea5535fd76bfc66df62d239bd94d0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:00:17 -0700 Subject: [PATCH 39/89] fix(mcp): reject scheme-only API key authorization payloads --- .../outbound_credentials/adapter.py | 5 ++++ .../outbound_credentials/test_adapter.py | 23 ++++++++++++++- .../mcp_server/test_mcp_server_manager.py | 5 +++- .../test_openapi_to_mcp_generator.py | 29 +++++++++++++++++++ 4 files changed, 60 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 5358878a248..fb9af25933a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -401,6 +401,11 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False + if auth_type == MCPAuth.api_key: + api_scheme: Final = value.split(None, 1)[0] + if api_scheme.lower() in ("bearer", "token", "apikey"): + api_credential: Final = strip_auth_scheme(value, api_scheme).strip() + return api_credential.lower() != api_scheme.lower() if auth_type in (MCPAuth.bearer_token, MCPAuth.token): scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" credential: Final = strip_auth_scheme(value, scheme).strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 2885fdaef95..4ada7c1763d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -7,6 +7,7 @@ maps each CredError onto its HTTP status. These pin the parity-critical mapping import base64 from types import SimpleNamespace +from typing import Final import pytest from fastapi import HTTPException @@ -20,7 +21,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_user_oauth_challenge, to_server_spec, to_subject, + validate_static_credential, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -34,10 +37,28 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( SharedKey, TokenExchangeConfig, ) -from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +@pytest.mark.parametrize("auth_type,header,value", [ + (MCPAuth.api_key, "Authorization", "Bearer fixture-key"), + (MCPAuth.api_key, "Authorization", "ApiKey fixture-key"), + (MCPAuth.api_key, "Authorization", "token fixture-key"), + (MCPAuth.api_key, "Authorization", "Bearer token"), + (MCPAuth.api_key, "Authorization", "opaque-key"), + (MCPAuth.api_key, "Authorization", "Custom Custom"), + (MCPAuth.api_key, "X-API-Key", "Bearer Bearer"), + (MCPAuth.api_key, "X-Custom", "ApiKey ApiKey"), + (MCPAuth.authorization, "Authorization", "Bearer Bearer"), +]) +def test_static_credential_preserves_supported_api_key_and_raw_headers( + auth_type: MCPAuthType, header: str, value: str, +) -> None: + result: Final = validate_static_credential(auth_type, {header: value}, upstream_token_header=header) + assert isinstance(result, Ok) + + def _server(**kwargs) -> MCPServer: return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 4c4c45162ca..9601f4dac4f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13767,7 +13767,10 @@ class TestProtectedCredentialPreparation: assert custom_slot is None or custom_slot not in request.headers @pytest.mark.asyncio - @pytest.mark.parametrize("value", ["", " ", "Bearer", "Basic", "token", "ApiKey"]) + @pytest.mark.parametrize("value", [ + "", " ", "Bearer", "Basic", "token", "ApiKey", + "Bearer Bearer", "ApiKey ApiKey", "token token", "bEaReR BEARER", "aPiKeY\tAPIKEY", + ]) async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: server: Final = MCPServer( server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 979199d0dc9..a7b7b0e9b44 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -40,6 +40,35 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +@pytest.mark.parametrize("value,accepted", [ + ("Bearer Bearer", False), ("ApiKey ApiKey", False), ("token token", False), + ("bEaReR BEARER", False), ("aPiKeY\tAPIKEY", False), + ("Bearer fixture-key", True), ("ApiKey fixture-key", True), ("token fixture-key", True), +]) +async def test_api_key_authorization_validates_payload_before_http( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, value: str, accepted: bool, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", auth_type=MCPAuth.api_key, + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(value) + try: + if accepted: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["authorization"] == value + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await tool() + assert exc.value.status_code == 500 + assert destination.call_count == 0 + finally: + _request_auth_header.reset(caller_token) + + @pytest.mark.asyncio @pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), From 7d42bc751debd3ffa7eaa6166d20f627cee0f067 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 10:05:13 -0700 Subject: [PATCH 40/89] docs(e2e): name all four rejection reasons in the counter note --- tests/e2e/PROVIDER_CACHE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index aca81f26c5a..e06b3c01653 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -46,7 +46,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all of one kind is a different problem from one whose rejections are nearly all of another, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. From 70ef8b24b6675faa14e7e3ec006d79debe032609 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:44:00 -0700 Subject: [PATCH 41/89] fix(mcp): reject bare schemes in raw authorization --- .../outbound_credentials/adapter.py | 2 +- .../outbound_credentials/test_adapter.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 35 ++++++++++++++++--- .../test_openapi_to_mcp_generator.py | 20 +++++++---- 4 files changed, 46 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index fb9af25933a..ba223f73b2d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -397,7 +397,7 @@ _STATIC_MODES: Final = frozenset( def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: if not value: return False - if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + if auth_type == MCPAuth.api_key and name != "authorization": return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 4ada7c1763d..78da9ff4d77 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -50,7 +50,7 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer (MCPAuth.api_key, "Authorization", "Custom Custom"), (MCPAuth.api_key, "X-API-Key", "Bearer Bearer"), (MCPAuth.api_key, "X-Custom", "ApiKey ApiKey"), - (MCPAuth.authorization, "Authorization", "Bearer Bearer"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), ]) def test_static_credential_preserves_supported_api_key_and_raw_headers( auth_type: MCPAuthType, header: str, value: str, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9601f4dac4f..186c46e1b37 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13577,19 +13577,46 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code in (401, 500) @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,slot", [(MCPAuth.api_key, "X-API-Key"), (MCPAuth.authorization, "Authorization")]) - async def test_raw_static_value_named_token_is_a_usable_credential(self, auth_type: MCPAuthType, slot: str) -> None: + @pytest.mark.parametrize("auth_type,slot,value", [ + (MCPAuth.api_key, "X-API-Key", "token"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), + (MCPAuth.authorization, "Authorization", "Bearer abc"), + (MCPAuth.authorization, "Authorization", "Custom abc"), + ]) + async def test_raw_static_credentials_are_forwarded_unchanged( + self, auth_type: MCPAuthType, slot: str, value: str, + ) -> None: server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, authentication_token="token") + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value) client = await MCPServerManager()._create_mcp_client(server) assert client._resolved_auth is not None request = httpx.Request("GET", server.url) flow = client._resolved_auth.auth_flow(request) try: - assert next(flow).headers[slot] == "token" + assert next(flow).headers[slot] == value finally: flow.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["Bearer", "basic", "token", "ApiKey", " bEaReR ", "\tTOKEN\t"]) + @pytest.mark.parametrize("source", ["configured", "caller", "forwarded"]) + async def test_raw_authorization_rejects_bare_schemes_before_dispatch( + self, respx_mock: MockRouter, value: str, source: str, + ) -> None: + server: Final = MCPServer( + server_id="raw-empty", name="raw-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.authorization, + authentication_token=value if source == "configured" else None, + ) + destination: Final = respx_mock.route().respond(200) + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=value if source == "caller" else None, + extra_headers={"Authorization": value} if source == "forwarded" else None, + ) + assert exc.value.status_code == 500 + assert destination.call_count == 0 + @pytest.mark.asyncio async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index a7b7b0e9b44..a9def20e75d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -41,17 +41,23 @@ GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp @pytest.mark.asyncio -@pytest.mark.parametrize("value,accepted", [ - ("Bearer Bearer", False), ("ApiKey ApiKey", False), ("token token", False), - ("bEaReR BEARER", False), ("aPiKeY\tAPIKEY", False), - ("Bearer fixture-key", True), ("ApiKey fixture-key", True), ("token fixture-key", True), +@pytest.mark.parametrize("auth_type,value,accepted", [ + (MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False), + (MCPAuth.api_key, "token token", False), (MCPAuth.api_key, "bEaReR BEARER", False), + (MCPAuth.api_key, "aPiKeY\tAPIKEY", False), (MCPAuth.api_key, "Bearer fixture-key", True), + (MCPAuth.api_key, "ApiKey fixture-key", True), (MCPAuth.api_key, "token fixture-key", True), + (MCPAuth.authorization, "Bearer", False), (MCPAuth.authorization, "basic", False), + (MCPAuth.authorization, "token", False), (MCPAuth.authorization, "ApiKey", False), + (MCPAuth.authorization, " bEaReR ", False), (MCPAuth.authorization, "\tTOKEN\t", False), + (MCPAuth.authorization, "opaque-secret-value", True), (MCPAuth.authorization, "Bearer abc", True), + (MCPAuth.authorization, "Custom abc", True), ]) -async def test_api_key_authorization_validates_payload_before_http( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, value: str, accepted: bool, +async def test_authorization_validates_credentials_before_http( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthType, value: str, accepted: bool, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", auth_type=MCPAuth.api_key, + "/echo", "get", {}, "https://upstream.example", auth_type=auth_type, ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") caller_token: Final = _request_auth_header.set(value) From 99aa9f76c8fc0a84969a975562dd04bbd3a7b160 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 18:51:22 +0000 Subject: [PATCH 42/89] fix(proxy): build failure headers immutably to keep LIT002 within budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/batches_endpoints/endpoints.py | 3 ++- litellm/proxy/common_utils/openai_error_payload.py | 7 +++++++ litellm/proxy/proxy_server.py | 4 ++-- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f37c06aea97..5d9ecddd4c2 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -7,6 +7,7 @@ import asyncio import os from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response @@ -731,7 +732,7 @@ async def list_batches( ) verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit) - data: dict = {} + data: Mapping[str, object] = MappingProxyType({}) try: if llm_router is None: raise HTTPException( diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index cbc8c78d4f9..202c61b620e 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -68,3 +68,10 @@ def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> Pr if litellm_call_id is not None: exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) return exc + + +def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]: + """``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name.""" + if headers is None: + return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id}) + return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f72cabf076..45beb8cc93c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -392,7 +392,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) from litellm.proxy.common_utils.openai_error_payload import ( - LITELLM_CALL_ID_HEADER, + headers_with_litellm_call_id, litellm_call_id_headers, with_litellm_call_id, ) @@ -11696,7 +11696,7 @@ async def audio_speech( raise HTTPException( status_code=e.status_code, detail=e.detail, - headers={LITELLM_CALL_ID_HEADER: litellm_call_id, **(e.headers or {})}, + headers=headers_with_litellm_call_id(e.headers, litellm_call_id), ) raise ProxyException( message=getattr(e, "message", f"{e}"), From 39acea0754e0dd5f291e8f99126b1e0b3f5505b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:42:15 -0700 Subject: [PATCH 43/89] feat(e2e): make Claude Code send the same bytes every build The compat cells drove the CLI with a fresh HOME per invocation and the pytest process's own working directory, and both reach the request body. The system prompt names a memory directory built from $CLAUDE_CONFIG_DIR/projects/, so a per-invocation config directory rewrote every body, and the CLI adds a git block for its working directory, so inheriting the checkout rewrote every body once per candidate. The device id churned for the same reason: the CLI mints it once and persists it in .claude.json, which we threw away each call. Nothing here was load-bearing. All three ride in metadata.user_id, whose job is abuse detection, not quota, caching or continuity. So pin the config directory and the working directory at fixed paths, seed the device id, and pin the session id. HOME stays fresh and empty per invocation, so the isolation is no weaker than before, and the CLI's own state no longer outlives the pod either. The working directory is deliberately not the checkout, so a model-directed Read now sees an empty directory rather than the repository. A pinned session id needs --no-session-persistence beside it: the CLI refuses a session id another live process holds, and the matrix runs its cells across xdist workers. Without the flag, six of eight concurrent invocations die on "Session ID is already in use". --- .../test_request_determinism.py | 143 ++++++++++++++++++ tests/e2e/claude_code/cli_driver.py | 57 +++++++ 2 files changed, 200 insertions(+) create mode 100644 tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py new file mode 100644 index 00000000000..09a181162d2 --- /dev/null +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -0,0 +1,143 @@ +"""The CLI must send the same request bytes from one build to the next. + +Markerless harness test: it drives the real `claude` binary against a local +stub instead of a proxy, so it carries no `e2e` marker. The binary is a +prerequisite of this whole suite, so a missing one is a failure rather than a +skip. + +Two builds differ in ways the driver does not control: a fresh pod, so no CLI +state survives, and a different candidate checked out at a different commit. +Both used to reach the request body, through the memory path the system prompt +names and through the git block the CLI adds for its working directory, so the +shared provider cache missed on every Claude Code cell. This replays those two +differences across a pair of invocations and holds the bytes equal. + +A pinned session id is what makes the second test necessary. The matrix runs +its cells across xdist workers, and the CLI refuses to start a session id that +another live process already holds, so pinning one without also opting out of +session persistence turns most of a parallel run red. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import threading +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import List, Tuple + +import pytest + +from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.rate_limiter import RateLimiter + +_STUB_REPLY = { + "id": "msg_stub", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 2}, +} + + +def _make_repo(root: Path, subject: str) -> Path: + root.mkdir(parents=True, exist_ok=True) + identity = {"NAME": "t", "EMAIL": "t@e2e"} + env = dict( + os.environ, + **{f"GIT_{role}_{key}": value for role in ("AUTHOR", "COMMITTER") for key, value in identity.items()}, + ) + (root / "file.txt").write_text(subject, encoding="utf-8") + for args in (["init", "-q"], ["add", "."], ["commit", "-q", "-m", subject]): + subprocess.run(["git", *args], cwd=root, env=env, check=True, capture_output=True) + return root + + +@pytest.fixture(name="captured") +def _captured() -> Tuple[str, List[bytes]]: + bodies: List[bytes] = [] + lock = threading.Lock() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + raw = self.rfile.read(int(self.headers.get("content-length") or 0)) + if "count_tokens" not in self.path: + with lock: + bodies.append(raw) + payload = json.dumps({"input_tokens": 10} if "count_tokens" in self.path else _STUB_REPLY).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *_args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", bodies + finally: + server.shutdown() + + +def test_two_builds_send_the_same_request_bytes(captured: Tuple[str, List[bytes]], tmp_path: Path) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + checkouts = (_make_repo(tmp_path / "build-1", "first"), _make_repo(tmp_path / "build-2", "second")) + origin = Path.cwd() + + sent = [] + for checkout in checkouts: + shutil.rmtree(Path(_stable_cli_state()[0]).parent, ignore_errors=True) + os.chdir(checkout) + try: + before = len(bodies) + run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ) + sent.append(bodies[before:]) + finally: + os.chdir(origin) + + assert sent[0], "the CLI sent no request to the stub, so there is nothing to compare" + assert sent[0] == sent[1] + + +def test_concurrent_cells_do_not_collide_on_the_pinned_session( + captured: Tuple[str, List[bytes]], tmp_path: Path +) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + + def one(_index: int) -> int: + return run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ).exit_code + + with ThreadPoolExecutor(max_workers=4) as pool: + codes = list(pool.map(one, range(4))) + + assert codes == [0, 0, 0, 0] + assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" + assert set(Counter(bodies).values()) == {4} diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 447e8cc0bbb..3fad87c7479 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -132,6 +132,57 @@ def _make_isolated_home() -> str: return tempfile.mkdtemp(prefix="claude-cli-home-") +_FIXED_CLI_USER_ID = "0" * 64 +_FIXED_CLI_SESSION_ID = "00000000-0000-4000-8000-000000000000" + + +def _seed_cli_identity(config_dir: str) -> None: + """Pin the device id the CLI would otherwise mint per config directory. + + It mints 32 random bytes on first run, writes them to `.claude.json` as + `userID`, and sends them in `metadata.user_id` forever after, so the value + is stable for exactly as long as that file lives. Pinning it, and the + session id passed beside it, costs nothing: both feed abuse detection + rather than quota, caching or continuity.""" + path = os.path.join(config_dir, ".claude.json") + try: + with open(path, encoding="utf-8") as handle: + if json.load(handle).get("userID") == _FIXED_CLI_USER_ID: + return + except (OSError, ValueError): + pass + staged = f"{path}.{os.getpid()}" + with open(staged, "w", encoding="utf-8") as handle: + json.dump({"userID": _FIXED_CLI_USER_ID}, handle) + os.replace(staged, path) + + +def _stable_cli_state() -> Tuple[str, str]: + """Config directory and working directory for the CLI, at fixed paths. + + Both reach the request body. The memory directory the system prompt + names is `$CLAUDE_CONFIG_DIR/projects//memory`, and a working + directory inside a git repository also contributes its branch and recent + commits. So a per-invocation config directory rewrites every body, and + inheriting the checkout rewrites every body once per candidate, which is + why the shared provider cache could never serve a Claude Code cell. + Pinning both makes the bodies repeatable across builds. + + This narrows what survives rather than widening it: HOME stays fresh and + empty per invocation, so the isolation `_make_isolated_home` describes is + unchanged, and the CLI's own state no longer outlives the pod either. The + working directory is deliberately not the checkout, so a model-directed + `Read` sees an empty directory instead of the repository. + """ + root = os.path.join(tempfile.gettempdir(), f"litellm-e2e-claude-{os.getuid()}") + config_dir = os.path.join(root, "config") + workspace = os.path.join(root, "workspace") + for path in (root, config_dir, workspace): + os.makedirs(path, mode=0o700, exist_ok=True) + _seed_cli_identity(config_dir) + return config_dir, workspace + + class ClaudeCLIError(RuntimeError): """Raised when the `claude` CLI cannot be invoked or returns a fatal error.""" @@ -222,6 +273,9 @@ def run_claude( "--verbose", "--model", model, + "--session-id", + _FIXED_CLI_SESSION_ID, + "--no-session-persistence", ] if extra_args: cmd.extend(extra_args) @@ -244,6 +298,8 @@ def run_claude( # regardless of how the subprocess exits. isolated_home = _make_isolated_home() env["HOME"] = isolated_home + config_dir, workspace = _stable_cli_state() + env["CLAUDE_CONFIG_DIR"] = config_dir if extra_env: env.update(extra_env) @@ -262,6 +318,7 @@ def run_claude( completed = run_fn( cmd, env=env, + cwd=workspace, input=stdin_input, capture_output=True, text=True, From 2481146727fe3b2613df96ccf8f568f5b0a3cc72 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:44:53 -0700 Subject: [PATCH 44/89] docs(e2e): say why the CLI-driving cells needed a driver fix, not a rule --- tests/e2e/PROVIDER_CACHE.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e06b3c01653..d37b37eeba7 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -16,6 +16,8 @@ Requests that differ only by their markers therefore share a canonical identity, Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to +A client that varies its own request between runs defeats that identity without breaking any rule, and the Claude Code compat cells did. The CLI sends a device id and a session id in `metadata.user_id`, and its system prompt names both its memory directory and its working directory, adding the branch and recent commits when that directory is a git repository. Driven with a fresh HOME and the checkout as its working directory, every cell sent different bytes every build. The fix belongs in the driver rather than here: `claude_code/cli_driver.py` pins the config directory, the working directory and both identifiers, which is why the cache needs no rule for any of it. Normalizing them instead would have hidden a real defect class, since a rule cannot tell a client's own churn from a value a test means to assert on + Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure From 9421b26bf6dcc98bee10e2cbbd45bf6ff1c23166 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:47:07 -0700 Subject: [PATCH 45/89] fix(e2e): stage the seeded device id per thread, not per process Build 232 took two compat cells red with a FileNotFoundError renaming `.claude.json.197` onto `.claude.json`. `run_claude_models_parallel` drives several models from one process, so a pid-suffixed staged name is shared between threads: one thread renamed the file the other was still writing, and the loser died on a path that no longer existed. mkstemp in the same directory gives a name that is unique per thread as well as per process, and the rename stays atomic. --- .../test_request_determinism.py | 19 ++++++++++++++++++- tests/e2e/claude_code/cli_driver.py | 11 ++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 09a181162d2..5046f35c73b 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -33,7 +33,7 @@ from typing import List, Tuple import pytest -from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter _STUB_REPLY = { @@ -141,3 +141,20 @@ def test_concurrent_cells_do_not_collide_on_the_pinned_session( assert codes == [0, 0, 0, 0] assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" assert set(Counter(bodies).values()) == {4} + + +def test_seeding_the_device_id_survives_threads_racing_on_the_same_directory(tmp_path: Path) -> None: + """`run_claude_models_parallel` drives several models from one process, so the + seed's staged file has to be unique per thread and not merely per process.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + seeded = config_dir / ".claude.json" + + for _round in range(20): + seeded.unlink(missing_ok=True) + with ThreadPoolExecutor(max_workers=16) as pool: + for outcome in [pool.submit(_seed_cli_identity, str(config_dir)) for _ in range(16)]: + outcome.result() + + assert json.loads(seeded.read_text(encoding="utf-8"))["userID"] == _FIXED_CLI_USER_ID + assert sorted(entry.name for entry in config_dir.iterdir()) == [".claude.json"] diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 3fad87c7479..a01d8ab3e7c 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -143,7 +143,12 @@ def _seed_cli_identity(config_dir: str) -> None: `userID`, and sends them in `metadata.user_id` forever after, so the value is stable for exactly as long as that file lives. Pinning it, and the session id passed beside it, costs nothing: both feed abuse detection - rather than quota, caching or continuity.""" + rather than quota, caching or continuity. + + The staged name has to be unique per *thread*, not per process: + `run_claude_models_parallel` drives several models from one process, so a + pid-suffixed name lets one thread rename the file another is still + writing, and the loser dies on a missing path.""" path = os.path.join(config_dir, ".claude.json") try: with open(path, encoding="utf-8") as handle: @@ -151,8 +156,8 @@ def _seed_cli_identity(config_dir: str) -> None: return except (OSError, ValueError): pass - staged = f"{path}.{os.getpid()}" - with open(staged, "w", encoding="utf-8") as handle: + handle_fd, staged = tempfile.mkstemp(dir=config_dir, prefix=".claude.json.") + with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: json.dump({"userID": _FIXED_CLI_USER_ID}, handle) os.replace(staged, path) From b96a80400441aa4073c837e21d9c542a0aa3814e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 20:55:32 +0000 Subject: [PATCH 46/89] fix(proxy): seed litellm_call_id into request data before parsing can fail The failure hook received data without the resolved id when body parsing or add_litellm_data_to_request raised, so proxy-only spend logging minted a fresh id that did not match the error log or the x-litellm-call-id header. The id is now part of the request data from the start and merged over the parsed body, which also removes the post-hoc in-place assignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/image_endpoints/endpoints.py | 5 +- .../pass_through_endpoints.py | 7 ++- litellm/proxy/proxy_server.py | 15 +++--- litellm/proxy/rerank_endpoints/endpoints.py | 5 +- .../proxy/image_endpoints/test_endpoints.py | 51 ++++++++++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 29 +++++++++++ 6 files changed, 92 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 30406bbcaae..5b90c0ff830 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -94,12 +94,12 @@ async def image_generation( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -110,7 +110,6 @@ async def image_generation( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if isinstance(model, str): reject_url_valued_destination("model", model) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e93b1232836..c27d4f17016 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -199,15 +199,15 @@ async def chat_completion_pass_through_endpoint( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: body: Final = await request.body() body_str: Final = body.decode() try: - data = ast.literal_eval(body_str) + data = ast.literal_eval(body_str) | data except Exception: - data = json.loads(body_str) + data = json.loads(body_str) | data data["adapter_id"] = adapter_id @@ -228,7 +228,6 @@ async def chat_completion_pass_through_endpoint( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id # override with user settings, these are params passed via cli if user_temperature: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 45beb8cc93c..2fb239c0c92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11471,12 +11471,12 @@ async def moderations( ``` """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11487,7 +11487,6 @@ async def moderations( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id data["model"] = ( general_settings.get("moderation_model", None) # server default @@ -11595,12 +11594,12 @@ async def audio_speech( https://platform.openai.com/docs/api-reference/audio/createSpeech """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11611,7 +11610,6 @@ async def audio_speech( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11730,12 +11728,12 @@ async def audio_transcriptions( https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly form_data: Final = await get_form_data(request) - data = {key: value for key, value in form_data.items() if key != "file"} + data = {key: value for key, value in form_data.items() if key != "file"} | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11746,7 +11744,6 @@ async def audio_transcriptions( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 4f5eb411e44..4f2daed15ed 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -58,11 +58,11 @@ async def rerank( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -73,7 +73,6 @@ async def rerank( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index d03832bf6d0..31b87530c94 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -1,7 +1,7 @@ import asyncio import copy import logging -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from types import SimpleNamespace from typing import Any, Dict @@ -279,3 +279,52 @@ async def test_failure_log_carries_the_callers_litellm_call_id( record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) assert record.litellm_call_id == call_id assert call_id in record.getMessage() + + +@pytest.mark.asyncio +async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIT-7836: when the request is rejected while it is still being prepared, the + failure hook must see the same litellm_call_id the response header answers with, + otherwise the spend row is stored under a freshly minted id nobody can look up.""" + call_id = "images-early-7836" + hook_request_data: list[Mapping[str, object]] = [] + + async def rejecting_add_litellm_data_to_request(**_: object) -> object: + raise HTTPException(status_code=400, detail={"error": "tag not allowed"}) + + async def fake_post_call_failure_hook(*, request_data: Mapping[str, object], **_: object) -> None: + hook_request_data.append(request_data) + + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", rejecting_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", + SimpleNamespace(post_call_failure_hook=fake_post_call_failure_hook), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version") + + body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk", "litellm_call_id": "from-the-body"}) + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/images/generations", + "headers": [(b"x-litellm-call-id", call_id.encode())], + }, + receive, + ) + + with pytest.raises(ProxyException) as raised: + await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.headers["x-litellm-call-id"] == call_id + assert [data["litellm_call_id"] for data in hook_request_data] == [call_id] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 14aea2bd020..a6e6a2100ca 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13046,6 +13046,35 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo assert call_id in record.getMessage() +@pytest.mark.asyncio +async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id(): + """LIT-7836: a body that fails to parse must still hand the failure hook the + litellm_call_id the response header answers with, so the spend row is findable.""" + from litellm.proxy._types import ProxyException + + call_id = "moderations-early-7836" + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": ') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + assert raised.value.headers["x-litellm-call-id"] == call_id + hook_request_data = fake_logging.post_call_failure_hook.await_args.kwargs["request_data"] + assert hook_request_data["litellm_call_id"] == call_id + + @pytest.mark.asyncio async def test_moderations_already_shaped_failure_answers_with_the_callers_litellm_call_id(): """LIT-7836: a ProxyException raised inside /v1/moderations is re-raised unwrapped but still From f10d95fb95371b1cebbec3ccd9fa3cce71fdc8ff Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:58:38 -0700 Subject: [PATCH 47/89] test(together_ai): drop the prefix-strip assertion, e2e covers it live tests/e2e/llm_translation/test_together_ai_e2e.py registers its model with the full registry key, so a slashed together_ai// goes through the prefix strip on every e2e run and an over-strip would fail against the real API. The unit assertion was a second copy of that. The roles check stays, since nothing in e2e exercises it. --- .../chat/test_together_ai_chat_transformation.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index a7347edb2c7..1df8c96fb50 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -1137,21 +1137,6 @@ def _together_chat_transport() -> tuple[HTTPHandler, list[httpx.Request]]: return client, captured_requests -def test_only_the_provider_prefix_is_stripped_from_a_slashed_model_name(): - client, captured_requests = _together_chat_transport() - - litellm.completion( - model=f"together_ai/{TOOL_CALLING_MODEL}", - messages=[{"role": "user", "content": "Hello!"}], - api_key="fake-key", - client=client, - ) - - assert "/" in TOOL_CALLING_MODEL - assert str(captured_requests[0].url) == "https://api.together.ai/v1/chat/completions" - assert json.loads(captured_requests[0].content)["model"] == TOOL_CALLING_MODEL - - def test_custom_role_wrappers_never_reach_the_request(): client, captured_requests = _together_chat_transport() messages = [{"role": "user", "content": "Hello!"}] From b4781317012c57c1bd7e186d1478b465316c6dfb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:20:49 -0700 Subject: [PATCH 48/89] test(together_ai): select the live model from the cost map `together_ai/openai/gpt-oss-20b` was hardcoded in two live tests and is no longer served, so both failed on a vendor catalog change rather than on anything litellm did. Both call sites now resolve the cheapest non-deprecated together_ai chat entry at runtime, filtered on the capabilities the tests actually exercise, mirroring what tests/e2e/llm_translation/test_together_ai_e2e.py already does. The selector lives in tests/_live_test_helpers.py so both lanes share one implementation. --- tests/_live_test_helpers.py | 36 +++++++++++++++++++++ tests/llm_translation/test_together_ai.py | 7 +++- tests/local_testing/test_text_completion.py | 7 +++- 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index a79b81e82c1..6b39f37921d 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -1,4 +1,7 @@ import os +from collections.abc import Mapping +from datetime import date +from typing import Any import pytest @@ -8,3 +11,36 @@ def _skip_live_prompt_caching_test(): pytest.skip("Live prompt-caching E2E tests are opt-in") if os.environ.get("CASSETTE_REDIS_URL"): pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay") + + +def cheapest_together_chat_model(*capability_flags: str) -> str: + import litellm + + today = date.today().isoformat() + + def qualifies(name: str, entry: Mapping[str, Any]) -> bool: + deprecation_date = entry.get("deprecation_date") + return ( + name.startswith("together_ai/") + and entry.get("litellm_provider") == "together_ai" + and entry.get("mode") == "chat" + and (deprecation_date is None or deprecation_date > today) + and (entry.get("input_cost_per_token") or 0.0) > 0 + and (entry.get("output_cost_per_token") or 0.0) > 0 + and all(bool(entry.get(flag)) for flag in capability_flags) + ) + + candidates = sorted( + ( + name + for name, entry in litellm.model_cost.items() + if isinstance(entry, Mapping) and qualifies(name, entry) + ), + key=lambda name: ( + litellm.model_cost[name].get("input_cost_per_token") or 0.0, + litellm.model_cost[name].get("output_cost_per_token") or 0.0, + name, + ), + ) + assert candidates, f"no live together_ai chat model in the cost map satisfies {capability_flags}" + return candidates[0] diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index fd7ad40ed11..7a49d46b528 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -3,6 +3,7 @@ Test TogetherAI LLM """ from base_llm_unit_tests import BaseLLMChatTest +from tests._live_test_helpers import cheapest_together_chat_model import json import os from datetime import datetime @@ -16,7 +17,11 @@ import pytest class TestTogetherAI(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: litellm.set_verbose = True - return {"model": "together_ai/openai/gpt-oss-20b"} + return { + "model": cheapest_together_chat_model( + "supports_function_calling", "supports_response_schema" + ) + } def test_tool_call_no_arguments(self, tool_call_no_arguments): """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833""" diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index a814ce6d303..b15037a2fcd 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -1,7 +1,11 @@ import asyncio import json +import os +import sys import traceback +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) + from dotenv import load_dotenv load_dotenv() @@ -12,6 +16,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm +from tests._live_test_helpers import cheapest_together_chat_model from litellm import ( RateLimitError, TextCompletionResponse, @@ -4030,7 +4035,7 @@ def test_async_text_completion_together_ai(): async def test_get_response(): try: response = await litellm.atext_completion( - model="together_ai/openai/gpt-oss-20b", + model=cheapest_together_chat_model(), prompt="good morning", max_tokens=10, ) From 09f3a5160a233556876460718607b440cc35aa37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:21:14 -0700 Subject: [PATCH 49/89] test(azure_ai): assert the bare deployment name reaches the native Responses endpoint --- .../test_azure_ai_responses_transformation.py | 86 ++++++++++++++----- 1 file changed, 66 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 1fe604d4a6b..925608a3c9b 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -192,21 +192,54 @@ def test_validate_environment_raises_without_credentials(): ) +NATIVE_RESPONSES_CASES = [ + ("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL, "gpt-5.6-luna-20260710154139"), + ( + "azure_ai/gpt-5.6-luna", + "https://res.services.ai.azure.com/models", + "https://res.services.ai.azure.com/openai/v1/responses", + "gpt-5.6-luna", + ), + ( + "azure_ai/gpt-5.6-sol", + "https://res.services.ai.azure.com", + "https://res.services.ai.azure.com/openai/v1/responses", + "gpt-5.6-sol", + ), + ( + "azure_ai/gpt-5.6-luna-20260710154139", + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + "gpt-5.6-luna-20260710154139", + ), + ( + "azure_ai/gpt-5.6-sol", + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + "gpt-5.6-sol", + ), +] + + +def _assert_native_responses_request(route, expected_url, expected_model): + request = route.calls.last.request + body = json.loads(request.content) + assert f"{request.url.scheme}://{request.url.host}{request.url.path}" == expected_url + assert request.headers["api-key"] == "fake-key" + assert body["model"] == expected_model + assert body["input"] == "What is the weather in SF?" + assert "messages" not in body + assert body["reasoning"] == {"effort": "high"} + assert body["tools"] == [WEATHER_TOOL] + + @pytest.mark.asyncio @respx.mock -@pytest.mark.parametrize( - "model,api_base,expected_url", - [ - ("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL), - ( - "azure_ai/gpt-5.6-luna", - "https://res.services.ai.azure.com/models", - "https://res.services.ai.azure.com/openai/v1/responses", - ), - ], -) -async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url): - route = respx.post(expected_url).mock(return_value=httpx.Response(200, json=_responses_payload("gpt-5.6-luna"))) +@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES) +async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url, expected_model): + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload(expected_model)) + ) await litellm.aresponses( model=model, @@ -217,13 +250,26 @@ async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, ap api_key="fake-key", ) - request = route.calls.last.request - body = json.loads(request.content) - assert request.headers["api-key"] == "fake-key" - assert body["input"] == "What is the weather in SF?" - assert "messages" not in body - assert body["reasoning"] == {"effort": "high"} - assert body["tools"] == [WEATHER_TOOL] + _assert_native_responses_request(route, expected_url, expected_model) + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES) +async def test_router_aresponses_sends_bare_deployment_name(model, api_base, expected_url, expected_model): + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload(expected_model)) + ) + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": model, "api_base": api_base, "api_key": "fake-key"}}], + num_retries=0, + ) + + await router.aresponses( + model="gpt-5.6", input="What is the weather in SF?", reasoning={"effort": "high"}, tools=[WEATHER_TOOL] + ) + + _assert_native_responses_request(route, expected_url, expected_model) @pytest.mark.asyncio From ba6c9fa61d87bd76634bdaba7ce4ef578771f923 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:26:16 -0700 Subject: [PATCH 50/89] fix(test): drop the redundant sys.path.insert CI runs these lanes as `python -m pytest` from the repo root, so the root is already on sys.path and `tests._live_test_helpers` imports without help. The insert only tripped the TQ003 test-quality budget. --- tests/local_testing/test_text_completion.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index b15037a2fcd..6808dfd768b 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -1,11 +1,7 @@ import asyncio import json -import os -import sys import traceback -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) - from dotenv import load_dotenv load_dotenv() From 515bf8c9d564731cfb285d47818db2fc081e1ba7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:32:08 -0700 Subject: [PATCH 51/89] refactor(test): validate cost-map entries into a typed model The selector read raw cost-map dicts as `Mapping[str, Any]`. It now validates each together_ai entry into a frozen Pydantic model and takes the two capabilities as keyword booleans, so nothing in the helper is coarsely typed or stringly addressed. --- tests/_live_test_helpers.py | 56 +++++++++++++++-------- tests/llm_translation/test_together_ai.py | 2 +- 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index 6b39f37921d..629f8ac9fdb 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -1,9 +1,8 @@ import os -from collections.abc import Mapping from datetime import date -from typing import Any import pytest +from pydantic import BaseModel, ConfigDict def _skip_live_prompt_caching_test(): @@ -13,34 +12,53 @@ def _skip_live_prompt_caching_test(): pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay") -def cheapest_together_chat_model(*capability_flags: str) -> str: + +class TogetherCostEntry(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str | None = None + mode: str | None = None + deprecation_date: str | None = None + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + supports_function_calling: bool | None = None + supports_response_schema: bool | None = None + + +def cheapest_together_chat_model( + *, function_calling: bool = False, response_schema: bool = False +) -> str: import litellm today = date.today().isoformat() - def qualifies(name: str, entry: Mapping[str, Any]) -> bool: - deprecation_date = entry.get("deprecation_date") + def qualifies(name: str, entry: TogetherCostEntry) -> bool: return ( name.startswith("together_ai/") - and entry.get("litellm_provider") == "together_ai" - and entry.get("mode") == "chat" - and (deprecation_date is None or deprecation_date > today) - and (entry.get("input_cost_per_token") or 0.0) > 0 - and (entry.get("output_cost_per_token") or 0.0) > 0 - and all(bool(entry.get(flag)) for flag in capability_flags) + and entry.litellm_provider == "together_ai" + and entry.mode == "chat" + and (entry.deprecation_date is None or entry.deprecation_date > today) + and (entry.input_cost_per_token or 0.0) > 0 + and (entry.output_cost_per_token or 0.0) > 0 + and (not function_calling or bool(entry.supports_function_calling)) + and (not response_schema or bool(entry.supports_response_schema)) ) + registry: dict[str, TogetherCostEntry] = { + name: TogetherCostEntry.model_validate(raw) + for name, raw in litellm.model_cost.items() + if isinstance(raw, dict) and name.startswith("together_ai/") + } candidates = sorted( - ( - name - for name, entry in litellm.model_cost.items() - if isinstance(entry, Mapping) and qualifies(name, entry) - ), + (name for name, entry in registry.items() if qualifies(name, entry)), key=lambda name: ( - litellm.model_cost[name].get("input_cost_per_token") or 0.0, - litellm.model_cost[name].get("output_cost_per_token") or 0.0, + registry[name].input_cost_per_token or 0.0, + registry[name].output_cost_per_token or 0.0, name, ), ) - assert candidates, f"no live together_ai chat model in the cost map satisfies {capability_flags}" + assert candidates, ( + "no live together_ai chat model in the cost map satisfies " + f"function_calling={function_calling} response_schema={response_schema}" + ) return candidates[0] diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index 7a49d46b528..0b4e9d3952c 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -19,7 +19,7 @@ class TestTogetherAI(BaseLLMChatTest): litellm.set_verbose = True return { "model": cheapest_together_chat_model( - "supports_function_calling", "supports_response_schema" + function_calling=True, response_schema=True ) } From 76c0f8db1d415a0307f4188a0ca992688ec3b44b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:01:08 -0700 Subject: [PATCH 52/89] chore(e2e): report the key components behind a mount that never converges Builds 232 and 233 held the Bedrock hit rate at 9% with the Claude Code driver already sending byte-identical requests and headers, so something between the proxy's ingress and the upstream still moves per build and the flat key cannot say what. Emit a digest per key component next to the counters: the test id, the method, the URL, each keyed header, the whole body, and one digest per top-level JSON body field. Values are digested, so no payload or credential reaches the artifact. Diffing two builds' artifacts names the field that moved. Diagnostic, to be removed once it has answered. --- tests/e2e/provider_cache.py | 63 +++++++++++++++++++++++++++++-- tests/e2e/provider_cache_redis.py | 4 ++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 444972ffce5..22967869c78 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -4,6 +4,8 @@ import base64 import hashlib import hmac import io +import json +import os import threading import time from collections.abc import Callable, Generator, Mapping @@ -400,6 +402,51 @@ def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: st return response +def component_digests( + test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> dict[str, str]: + """Per-component digests of everything the key covers. + + A mount whose corpus never converges is a mount where one of these moves + between builds, and the flat key cannot say which. Values are digested, so + no payload or credential is written, and a JSON body contributes one digest + per top-level field so the field that moved can be named.""" + parts: dict[str, str] = { # rebind-ok: a report assembled from three differently shaped sources + "test_key": test_key, + "method": method, + "url": short_digest(canonical_text(url).encode()), + } + for name, value in sorted(headers.items()): + parts[f"header:{name.lower()}"] = short_digest(value.encode()) + canonical: Final = b"" if body is None else canonical_body(body) + parts["body"] = short_digest(canonical) + try: + parsed: Final = JSON_VALUE.validate_json(canonical) + except ValidationError: + return parts + if isinstance(parsed, dict): + for name, value in sorted(parsed.items()): + parts[f"body:{name}"] = short_digest(json.dumps(value, sort_keys=True).encode()) + return parts + + +def short_digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest()[:16] + + +@dataclass(slots=True) +class KeyProbe: + """Every keyed request's components, when a metrics directory is configured.""" + + rows: tuple[tuple[tuple[str, str], ...], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def observe(self, mount: str, outcome: str, parts: Mapping[str, str]) -> None: + row: Final = tuple({"mount": mount, "outcome": outcome, **parts}.items()) + with self.lock: + self.rows = (*self.rows, row) + + @dataclass(slots=True) class CacheCounters: counts: tuple[tuple[str, int], ...] = () @@ -464,6 +511,7 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + probe: KeyProbe = field(default_factory=KeyProbe) slots: SlotCounter = field(default_factory=SlotCounter) policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 @@ -481,6 +529,14 @@ class CacheEdge: self.counters.increment(name) self.counters.increment(f"mount:{mount}:{name}") + def record_key( + self, mount: str, outcome: str, test_key: str, method: str, url: str, + headers: Mapping[str, str], body: bytes | None, + ) -> None: + if not os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR"): + return + self.probe.observe(mount, outcome, component_digests(test_key, method, url, headers, body)) + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy @@ -511,20 +567,21 @@ class CacheEdge: if isinstance(prepared, NetworkError): self.reject(mount, UNREACHABLE) return prepared - identity: Final = request_identity( - self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, - ) + keyed_headers: Final = self.keyed(mount, prepared.headers) + identity: Final = request_identity(self.secret, test_key, method, url, keyed_headers, body) key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: self.count(mount, "hits") + self.record_key(mount, "hit", test_key, method, url, keyed_headers, body) return StreamHead(response.status_code, response.headers, response_steps(response)) self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found self.count(mount, "misses") + self.record_key(mount, "miss", test_key, method, url, keyed_headers, body) if isinstance(capture_slot, CacheUnavailable): self.count(mount, "cache_errors") self.count(mount, "upstream_attempts") diff --git a/tests/e2e/provider_cache_redis.py b/tests/e2e/provider_cache_redis.py index be4e31b2c49..2c7419cfc0f 100644 --- a/tests/e2e/provider_cache_redis.py +++ b/tests/e2e/provider_cache_redis.py @@ -134,6 +134,10 @@ def write_metrics(cache: CacheEdge) -> None: root: Final = Path(directory) root.mkdir(parents=True, exist_ok=True) (root / f"{os.getpid()}.json").write_text(report + "\n") + if cache.probe.rows: + (root / f"keys-{os.getpid()}.json").write_text( + json.dumps([dict(row) for row in cache.probe.rows]) + "\n" + ) except OSError: logging.getLogger(__name__).warning("provider cache metrics artifact unavailable") logging.getLogger(__name__).info("%s", report) From ea5887d4bb5eec62dfab82b2f71183e210cf75b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:10:48 -0700 Subject: [PATCH 53/89] test: cover proxy_admin_viewer sessions in the model_group/info admin regression test --- .../proxy/proxy_server/test_routes_model_info.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index a076b3593b1..2c101156c6c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -340,10 +340,13 @@ def model_group_info_router(monkeypatch): return router -def test_model_group_info_proxy_admin_ignores_key_model_restriction(client, auth_as, model_group_info_router): +@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) +def test_model_group_info_proxy_admin_ignores_key_model_restriction( + client, auth_as, model_group_info_router, admin_role +): from litellm.proxy._types import LitellmUserRoles - with auth_as(LitellmUserRoles.PROXY_ADMIN, models=["no-default-models"]): + with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]): response = client.get("/model_group/info") assert response.status_code == 200 From a9fc6d255b5bd82220f3c2265eaf568dcf179423 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:15:42 -0700 Subject: [PATCH 54/89] fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge Six Amazon Nova entries define cache_read_input_token_cost twice, which is what a clean text merge of two branches that both added the field looks like. JSON parsers keep the last occurrence, so this turned test_price_map_has_no_duplicate_keys red on every open PR's merge commit, including this one, which touches neither file. Both occurrences in all six entries carry the same value, so dropping the later one leaves every parsed price identical. Same change as #41496, carried here so this branch is not blocked on it. Identical deletions, so the two merge cleanly in either order. --- ...model_prices_and_context_window_backup.json | 18 ++++++------------ model_prices_and_context_window.json | 18 ++++++------------ 2 files changed, 12 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "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 b1c38d0350a..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From b9751a38ab0a21f56703569b18b3db084a6fa1ae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:36:54 -0700 Subject: [PATCH 55/89] Revert "fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge" This reverts commit a9fc6d255b5bd82220f3c2265eaf568dcf179423. --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++------ model_prices_and_context_window.json | 18 ++++++++++++------ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 92e5b1c4ff7..b1c38d0350a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "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 92e5b1c4ff7..b1c38d0350a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From 9b77b5c2cb5ec8bbf58738d279e266e4f8162212 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 18:08:09 -0700 Subject: [PATCH 56/89] feat(proxy): let proxy admins choose which team fields team admins may edit Team admins could never reach POST /team/update: the route gate answered 401 before the handler's team-admin branch ran. This moves /team/update into the self-managed routes, resolves proxy admin, org admin or team admin inside the handler, and filters team admins through a new proxy-wide UI setting, team_admin_editable_team_fields. The setting is an allow-list of team fields. Empty means team admins cannot edit team settings and get a 403 pointing at the proxy admin, and changing a field outside the list fails 403 naming that field. Only values that differ from what is stored count, since the dashboard resends the whole form. The registry of fields the setting accepts ships empty on purpose. Each field lands in its own follow-up PR with its value diff and dashboard wiring. The Admin UI gains a "Team admin editable fields" section under Settings > UI and a toast on the team page while editing is disabled. Refs LIT-5722 Claude-Session: https://claude.ai/code/session_01A6SkwJdfZUmkzfUkrEkqX8 --- litellm/proxy/_types.py | 3 + .../team_admin_field_permissions.py | 179 ++++++++++++++ .../management_endpoints/team_endpoints.py | 84 +++++-- .../proxy_setting_endpoints.py | 35 +++ .../management/test_team_update.py | 63 ++--- .../proxy/auth/test_route_checks.py | 83 ++++--- .../test_team_admin_field_permissions.py | 125 ++++++++++ .../test_team_endpoints.py | 234 +++++++++++++++++- .../test_proxy_setting_endpoints.py | 108 ++++++++ .../TeamAdminEditableFieldsSettings.test.tsx | 91 +++++++ .../TeamAdminEditableFieldsSettings.tsx | 64 +++++ .../UISettings/UISettings.test.tsx | 44 ++++ .../AdminSettings/UISettings/UISettings.tsx | 26 ++ .../src/components/team/TeamInfo.test.tsx | 65 +++++ .../src/components/team/TeamInfo.tsx | 27 +- .../team/teamAdminEditAccess.test.ts | 69 ++++++ .../components/team/teamAdminEditAccess.ts | 38 +++ 17 files changed, 1244 insertions(+), 94 deletions(-) create mode 100644 litellm/proxy/management_endpoints/team_admin_field_permissions.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx create mode 100644 ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 14e3635f079..94cfd090dcc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum): ) self_managed_routes = [ + # update_team resolves proxy/org/team admin itself and filters team admins + # through the team_admin_editable_team_fields setting + "/team/update", "/team/member_add", "/team/member_delete", "/management/v1/teams/{team_id}/members/bulk_delete", diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py new file mode 100644 index 00000000000..9ccacf9c54f --- /dev/null +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -0,0 +1,179 @@ +"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, assert_never + +from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + UpdateTeamRequest, +) + +TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields" + +# TODO(LIT-5722): stays empty until each field's value-diff and dashboard wiring lands, one field per PR +SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset() + +_FIELD_LIST: Final = TypeAdapter(list[str]) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset( + (*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium) +) +_SYSTEM_MANAGED_METADATA_KEYS: Final[frozenset[str]] = frozenset({"team_member_budget_id"}) +_NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"}) +_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields" + + +@dataclass(frozen=True, slots=True) +class TeamAdminEditAllowed: + kind: Literal["allowed"] = "allowed" + + +@dataclass(frozen=True, slots=True) +class TeamAdminEditingDisabled: + kind: Literal["disabled"] = "disabled" + + +@dataclass(frozen=True, slots=True) +class TeamAdminFieldNotPermitted: + field: str + kind: Literal["field_not_permitted"] = "field_not_permitted" + + +TeamAdminEditVerdict: TypeAlias = TeamAdminEditAllowed | TeamAdminEditingDisabled | TeamAdminFieldNotPermitted + + +def resolve_team_admin_editable_fields( + general_settings: Mapping[str, object], + supported: frozenset[str], +) -> frozenset[str]: + raw: Final = general_settings.get(TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING) + if raw is None: + return frozenset() + try: + configured: Final = frozenset(_FIELD_LIST.validate_python(raw)) + except ValidationError: + verbose_proxy_logger.warning( + "%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw + ) + return frozenset() + unsupported: Final = configured - supported + if unsupported: + verbose_proxy_logger.warning( + "%s ignores unsupported field(s) %s; supported: %s", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, + sorted(unsupported), + sorted(supported), + ) + return configured & supported + + +def _as_object(value: object) -> Mapping[str, object]: + try: + return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) + except ValidationError: + return _EMPTY + + +def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]: + return _as_object(existing.get("metadata")) + + +def _submitted_metadata( + data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] +) -> Mapping[str, object]: + """Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over.""" + base: Final = ( + _as_object(submitted.get("metadata")) if "metadata" in data.model_fields_set else _stored_metadata(existing) + ) + folded: Final = data.model_fields_set & _METADATA_FOLDED_FIELDS + return MappingProxyType({key: submitted[key] if key in folded else base[key] for key in base.keys() | folded}) + + +def _metadata_changes( + data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] +) -> frozenset[str]: + merged: Final = _submitted_metadata(data, submitted, existing) + stored: Final = _stored_metadata(existing) + return frozenset( + key if key in _METADATA_FOLDED_FIELDS else "metadata" + for key in (merged.keys() | stored.keys()) - _SYSTEM_MANAGED_METADATA_KEYS + if merged.get(key) != stored.get(key) + ) + + +def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]: + table: Final = existing_row.litellm_model_table + return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY + + +def _column_changed( + field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable +) -> bool: + if field == "model_aliases": + return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row) + if field in LiteLLM_TeamTable.model_fields: + return submitted.get(field) != existing.get(field) + return True + + +def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]: + """Logical field names whose stored value the request would change. + + Request and stored row are compared as JSON values so both sides share one representation. Fields the + server folds into metadata are attributed to their own name whether they arrive top-level or inside + ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored + counterpart on the team row count as changed whenever they are sent. + """ + submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True)) + existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json()) + column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS + column_changes: Final = frozenset( + field for field in column_fields if _column_changed(field, submitted, existing, existing_row) + ) + return column_changes | _metadata_changes(data, submitted, existing) + + +def team_admin_edit_verdict( + data: UpdateTeamRequest, + existing: LiteLLM_TeamTable, + permitted: frozenset[str], +) -> TeamAdminEditVerdict: + if not permitted: + return TeamAdminEditingDisabled() + blocked: Final = sorted(changed_team_fields(data, existing) - permitted) + if blocked: + return TeamAdminFieldNotPermitted(field=blocked[0]) + return TeamAdminEditAllowed() + + +def raise_for_team_admin_edit_verdict(verdict: TeamAdminEditVerdict) -> None: + match verdict: + case TeamAdminEditAllowed(): + return + case TeamAdminEditingDisabled(): + raise HTTPException( + status_code=403, + detail=( + "Team admins on this proxy cannot edit team settings. " + f"Ask a proxy admin to enable fields under {_SETTINGS_LOCATION}." + ), + ) + case TeamAdminFieldNotPermitted(field=field): + raise HTTPException( + status_code=403, + detail=( + f"Team admins on this proxy do not have permission to update '{field}'. " + f"Ask a proxy admin to add it under {_SETTINGS_LOCATION}." + ), + ) + case _: + assert_never(verdict) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 23b02c22f24..d88d6d13720 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -18,7 +18,18 @@ from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast +from typing import ( + TYPE_CHECKING, + Annotated, + Final, + Literal, + NamedTuple, + NoReturn, + Protocol, + TypeAlias, + TypeVar, + cast, +) import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -122,6 +133,12 @@ from litellm.proxy.management_endpoints.router_weights import validate_router_se from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + raise_for_team_admin_edit_verdict, + resolve_team_admin_editable_fields, + team_admin_edit_verdict, +) from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, AccessGroupSyncTx, @@ -439,32 +456,43 @@ async def _refresh_cached_team( ) -async def _can_manage_team( +TeamAccessRole: TypeAlias = Literal["proxy_admin", "org_admin", "team_admin"] + + +def _raise_team_access_denied() -> NoReturn: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have access to this team", + ) + + +async def _resolve_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, -) -> bool: - """True for a proxy admin, an admin of this team, or an org admin for the team's organization.""" +) -> TeamAccessRole | None: + """Strongest role the caller holds over ``team_obj``, or None when they hold none. + + Org admin outranks team admin so a caller holding both keeps unrestricted edits. + """ if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - return True + return "proxy_admin" + + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): + return "org_admin" if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): - return True + return "team_admin" - return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + return None async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Raise HTTPException(403) unless the caller can manage the given team.""" - if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict): - return - - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="You do not have access to this team", - ) + """Raise 403 unless the caller is a proxy admin, an org admin for the team's org, or a team admin.""" + if await _resolve_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) is None: + _raise_team_access_denied() class TeamMemberBudgetHandler: @@ -2096,6 +2124,7 @@ async def update_team( try: from litellm.proxy.management_helpers.audit_logs import is_audit_logging_enabled from litellm.proxy.proxy_server import ( + general_settings, litellm_proxy_admin_name, llm_router, premium_user, @@ -2144,16 +2173,29 @@ async def update_team( ) if existing_team_row is None: + # Non-proxy-admins get the same 403 as an access denial so /team/update + # cannot be used to probe which team ids exist + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + _raise_team_access_denied() raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - # Verify caller has access to manage this team - await _verify_team_access( - team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()), - user_api_key_dict=user_api_key_dict, - ) + existing_team: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()) + access_role: Final = await _resolve_team_access(team_obj=existing_team, user_api_key_dict=user_api_key_dict) + if access_role is None: + _raise_team_access_denied() + if access_role == "team_admin": + raise_for_team_admin_edit_verdict( + team_admin_edit_verdict( + data=data, + existing=existing_team, + permitted=resolve_team_admin_editable_fields( + general_settings, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + ), + ) + ) await validate_router_settings_weights( data.router_settings, @@ -4585,7 +4627,7 @@ async def team_info( await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table) organization_models: Final[list[str] | None] = ( _parent_organization_models(team_info) - if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict) + if await _resolve_team_access(team_obj=team_table, user_api_key_dict=user_api_key_dict) is not None else None ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index c12d071dd36..c69225c55d6 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -29,6 +29,10 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, +) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository @@ -212,6 +216,9 @@ class UIThemeSettingsResponse(SettingsResponse): """Response model for UI theme settings""" +_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)) + + class UISettings(BaseModel): """Configuration for UI-specific flags""" @@ -304,6 +311,18 @@ class UISettings(BaseModel): description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) + team_admin_editable_team_fields: Sequence[str] = Field( + default=(), + description=( + "Team settings fields a team admin may change on the teams they administer. " + "Empty means team admins cannot edit team settings at all. " + "Proxy admins and org admins are not affected." + ), + json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict + "items": {"type": "string", "enum": [*_TEAM_ADMIN_FIELD_ENUM]}, # mutable-ok: nested in the dict above + }, + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -326,6 +345,7 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = { "disable_custom_api_keys", "disable_key_generate_for_org_admin", "enable_chat_ui", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, } ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution" @@ -360,6 +380,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", "disable_key_generate_for_org_admin", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, ] # Extension point: packages outside OSS (e.g. litellm_enterprise) can @@ -1571,6 +1592,20 @@ async def update_ui_settings( except ValidationError as e: raise HTTPException(status_code=422, detail=e.errors()) + unsupported_team_fields: Final = sorted( + frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + ) + if unsupported_team_fields: + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization + "error": ( + f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. " + f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}." + ) + }, + ) + # Only include fields the caller actually sent (not Pydantic defaults). settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) diff --git a/tests/proxy_behavior/management/test_team_update.py b/tests/proxy_behavior/management/test_team_update.py index 23ea89fa74d..d6273e64132 100644 --- a/tests/proxy_behavior/management/test_team_update.py +++ b/tests/proxy_behavior/management/test_team_update.py @@ -9,31 +9,31 @@ pytestmark = pytest.mark.asyncio(loop_scope="session") # POST /team/update — actor x team-shape matrix (shapes built by _seed_target). -# Each request carries the team's own organization_id so a non-proxy-admin can -# reach the org-scoped branch of the route-permission gate (401 on denial), -# which fronts the handler's _verify_team_access. Only PROXY_ADMIN and an -# ORG_ADMIN of the team's org pass: an internal_user team admin is filtered by -# the route gate before _verify_team_access's team-admin branch is reached. +# The route is self-managed (LIT-5722), so every authenticated caller reaches +# update_team and denials are the handler's 403, never the route gate's 401. +# Only PROXY_ADMIN and an ORG_ADMIN of the team's org pass: a team admin is +# admitted by _resolve_team_access but then refused because no team field is +# enabled for team admins (team_admin_editable_team_fields ships empty). MARKER_ALIAS = "behavior-pin-update-marker-alias" _MATRIX = [ ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), - ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401), - ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401), - ("alpha/owner", Actor.OWNER, "alpha", 401), - ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401), - ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401), - ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401), - ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 403), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), - ("beta/org_admin", Actor.ORG_ADMIN, "beta", 401), - ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 401), - ("beta/internal_user", Actor.INTERNAL_USER, "beta", 401), - ("beta/owner", Actor.OWNER, "beta", 401), - ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 401), - ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 401), - ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 401), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), ] @@ -110,8 +110,9 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context( ): """With no organization_id in the body the route gate resolves the target team's org from team_id, so an org admin of the team's own org is allowed - (200), same as PROXY_ADMIN. A team admin of that same team stays denied - (401): the resolution grants org admins access, not team admins.""" + (200), same as PROXY_ADMIN. A team admin of that same team reaches the + handler but is refused (403) until a proxy admin enables fields for team + admins, and the response says so.""" await _seed_target(prisma, world, "alpha", scratch.prefix) allowed_org_admin = await proxy_client.post( @@ -133,21 +134,25 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context( headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"}, json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS}, ) - assert denied_team_admin.status_code == 401, denied_team_admin.text + assert denied_team_admin.status_code == 403, denied_team_admin.text + assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text + assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text # Relocation gate — moving a team to a different org. The scratch team starts # in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses; -# ORG_B_ADMIN clears the route gate (dest-org admin) but fails -# _verify_team_access on the source team (403); the rest fail the route gate -# (401). The relocation-*allowed* branch (caller is org admin of both orgs) is -# covered by test_team_update_org_relocation_allowed_for_dual_org_admin below. +# ORG_B_ADMIN reaches the handler but holds no role on the source team (403); +# ORG_ADMIN holds the source team but not the destination org (403 from the +# relocation gate); the team admin is refused by the empty field allow-list and +# the internal user holds no role at all (403). The relocation-*allowed* branch +# (caller is org admin of both orgs) is covered by +# test_team_update_org_relocation_allowed_for_dual_org_admin below. _RELOCATION = [ ("proxy_admin", Actor.PROXY_ADMIN, 200), ("org_b_admin", Actor.ORG_B_ADMIN, 403), - ("org_admin", Actor.ORG_ADMIN, 401), - ("team_admin", Actor.TEAM_ADMIN, 401), - ("internal_user", Actor.INTERNAL_USER, 401), + ("org_admin", Actor.ORG_ADMIN, 403), + ("team_admin", Actor.TEAM_ADMIN, 403), + ("internal_user", Actor.INTERNAL_USER, 403), ] diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 806c55d51ce..ed319190a6b 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2892,45 +2892,53 @@ def test_team_update_gate_allows_org_admin_with_resolved_org(): ) -def test_team_update_gate_rejects_without_org_context(): - """Without organization_id (i.e. resolution found no org, or a non-org-admin), - the gate still rejects /team/update — the fix adds no blanket allow. Guards - against re-widening the route (e.g. dropping it into self_managed_routes).""" +def test_team_update_gate_admits_internal_user_without_org_context(): + """/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated + caller and update_team resolves proxy, org or team admin itself, then filters team admins + through the team_admin_editable_team_fields setting. Before that the gate 401'd every + team admin, which left the handler's team-admin branch unreachable.""" + from litellm.proxy._types import LiteLLMRoutes + + assert "/team/update" in LiteLLMRoutes.self_managed_routes.value + + user_obj = LiteLLM_UserTable( + user_id="team-admin-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + organization_memberships=None, + ) + valid_token = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/team/update", + request=request, + valid_token=valid_token, + request_data={"team_id": "team-1", "max_budget": 42}, + ) + + +def test_team_update_gate_defers_cross_org_admin_to_the_handler(): # test-quality-ok: the gate's only success signal is not raising; the handler's 403 it defers to is pinned in test_team_endpoints + """An org admin of a DIFFERENT org clears the coarse gate like any internal user; + update_team's _resolve_team_access finds no role on the team and 403s (pinned in + test_team_endpoints), so there is still no cross-org escalation.""" user_obj = _make_org_admin_user("org-1") valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) request = MagicMock(spec=Request) request.method = "POST" request.query_params = {} - with pytest.raises(Exception, match="Only proxy admin can be used to generate"): - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=LitellmUserRoles.INTERNAL_USER.value, - route="/team/update", - request=request, - valid_token=valid_token, - request_data={"team_id": "team-1", "max_budget": 42}, - ) - - -def test_team_update_gate_rejects_cross_org_admin_with_resolved_org(): - """Even after the target team's org is resolved, an org admin of a DIFFERENT - org is rejected at the gate (no cross-org escalation).""" - user_obj = _make_org_admin_user("org-1") - valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) - request = MagicMock(spec=Request) - request.method = "POST" - request.query_params = {} - - with pytest.raises(Exception, match="Only proxy admin can be used to generate"): - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=LitellmUserRoles.INTERNAL_USER.value, - route="/team/update", - request=request, - valid_token=valid_token, - request_data={"team_id": "team-1", "organization_id": "org-2"}, - ) + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/team/update", + request=request, + valid_token=valid_token, + request_data={"team_id": "team-1", "organization_id": "org-2"}, + ) # ── PATCH /team/{team_id}: same org-context + role reach as POST /team/update ── @@ -2993,10 +3001,11 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_has_same_reach_as_team_update(): - """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but - NOT by regular internal users or the role-agnostic self_managed_routes — the - latter would open /team/new (the collision footgun) to any authenticated user.""" +def test_patch_team_route_stays_out_of_self_managed_routes(): + """Unlike POST /team/update, PATCH /team/{team_id} cannot be self-managed: its + template also matches /team/new (the collision footgun), so it stays reachable by + org admins (org_admin_allowed_routes) and proxy admins only, never by regular + internal users or through the role-agnostic self_managed_routes.""" from litellm.proxy._types import LiteLLMRoutes assert RouteChecks.check_route_access( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py new file mode 100644 index 00000000000..91479921c61 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -0,0 +1,125 @@ +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateTeamRequest +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + TeamAdminEditAllowed, + TeamAdminEditingDisabled, + TeamAdminFieldNotPermitted, + changed_team_fields, + raise_for_team_admin_edit_verdict, + resolve_team_admin_editable_fields, + team_admin_edit_verdict, +) + +_SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"}) + + +def _team(**overrides): + return LiteLLM_TeamTable(team_id="team-1", **overrides) + + +class TestResolveTeamAdminEditableFields: + def test_missing_setting_means_nothing_editable(self): + assert resolve_team_admin_editable_fields({}, _SUPPORTED) == frozenset() + + def test_keeps_only_supported_names(self): + configured = {"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]} + assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"}) + + @pytest.mark.parametrize("raw", ["tpm_limit", 7, {"tpm_limit": True}, [1, 2]]) + def test_malformed_setting_fails_closed(self, raw): + assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset() + + +class TestChangedTeamFields: + def test_team_id_alone_changes_nothing(self): + assert changed_team_fields(UpdateTeamRequest(team_id="team-1"), _team()) == frozenset() + + def test_column_echoing_stored_value_is_not_a_change(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=5, team_alias="alpha", max_budget=None) + assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset() + + def test_column_with_different_value_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") + assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset({"tpm_limit"}) + + def test_explicit_null_clearing_a_stored_column_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", max_budget=None) + assert changed_team_fields(data, _team(max_budget=30.0)) == frozenset({"max_budget"}) + + def test_folded_field_sent_top_level_is_named_not_metadata(self): + data = UpdateTeamRequest(team_id="team-1", guardrails=["b"]) + assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"}) + + def test_folded_field_sent_inside_metadata_is_named_not_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}) + assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"}) + + def test_custom_metadata_key_change_is_attributed_to_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["a"], "cost_center": "b"}) + existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"}) + assert changed_team_fields(data, existing) == frozenset({"metadata"}) + + def test_metadata_echo_with_top_level_override_only_names_the_override(self): + data = UpdateTeamRequest(team_id="team-1", guardrails=["b"], metadata={"guardrails": ["a"], "cost_center": "a"}) + existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"}) + assert changed_team_fields(data, existing) == frozenset({"guardrails"}) + + def test_dropping_a_stored_key_from_submitted_metadata_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"}) + existing = _team(metadata={"cost_center": "a", "tags": ["x"], "logging": [{"callback": "langfuse"}]}) + assert changed_team_fields(data, existing) == frozenset({"tags", "logging"}) + + def test_server_managed_metadata_key_is_ignored(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"}) + existing = _team(metadata={"cost_center": "a", "team_member_budget_id": "budget-1"}) + assert changed_team_fields(data, existing) == frozenset() + + def test_model_aliases_compare_against_the_model_table(self): + table = LiteLLM_ModelTable(model_aliases='{"fast": "gpt-4o-mini"}', created_by="a", updated_by="a") + same = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o-mini"}) + different = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o"}) + assert changed_team_fields(same, _team(litellm_model_table=table)) == frozenset() + assert changed_team_fields(different, _team(litellm_model_table=table)) == frozenset({"model_aliases"}) + + def test_empty_model_aliases_against_no_model_table_is_not_a_change(self): + assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset() + + def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self): + data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0) + assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"}) + + +class TestTeamAdminEditVerdict: + def test_no_permitted_fields_disables_editing_even_for_a_no_op(self): + verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset()) + assert verdict == TeamAdminEditingDisabled() + + def test_changes_within_permitted_fields_are_allowed(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") + verdict = team_admin_edit_verdict(data, _team(team_alias="alpha"), frozenset({"tpm_limit"})) + assert verdict == TeamAdminEditAllowed() + + def test_first_blocked_field_in_sorted_order_is_reported(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True) + verdict = team_admin_edit_verdict(data, _team(), frozenset({"tpm_limit"})) + assert verdict == TeamAdminFieldNotPermitted(field="blocked") + + +class TestRaiseForTeamAdminEditVerdict: + def test_allowed_does_not_raise(self): + assert raise_for_team_admin_edit_verdict(TeamAdminEditAllowed()) is None + + def test_disabled_is_a_403_pointing_at_the_proxy_admin(self): + with pytest.raises(HTTPException) as exc: + raise_for_team_admin_edit_verdict(TeamAdminEditingDisabled()) + assert exc.value.status_code == 403 + assert "cannot edit team settings" in exc.value.detail + assert "Settings > UI > Team admin editable fields" in exc.value.detail + + def test_field_not_permitted_is_a_403_naming_the_field(self): + with pytest.raises(HTTPException) as exc: + raise_for_team_admin_edit_verdict(TeamAdminFieldNotPermitted(field="blocked")) + assert exc.value.status_code == 403 + assert "'blocked'" in exc.value.detail 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 6478b18e553..733f6ddc36f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1,6 +1,6 @@ import asyncio import json -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import Final, Optional, cast @@ -76,6 +76,31 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( client = TestClient(app) +@contextmanager +def _team_admin_may_edit(*fields: str): + """Let team admins change ``fields`` on /team/update for the duration of the block. + + The registry ships empty (LIT-5722 adds fields one PR at a time), so tests that exercise the + gates layered underneath the allow-list widen it here instead of asserting the early 403.""" + with ( + patch( # test-quality-ok: the registry is a module constant update_team reads directly; no seam to inject + "litellm.proxy.management_endpoints.team_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset(fields), + ), + patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": list(fields)}), # test-quality-ok: update_team reads general_settings as a proxy_server module global + ): + yield + + +def _not_org_admin(): + """update_team asks whether the caller administers the team's org before it settles for team admin; + a MagicMock prisma cannot answer that lookup, so pin it to False.""" + return patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=False), + ) + + def _wire_team_create_tx(prisma_client): """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, so a mocked client has to hand its team table back out of `db.tx()`. @@ -6393,6 +6418,7 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin(): dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6549,6 +6575,7 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin(): dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6618,6 +6645,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6712,6 +6740,7 @@ async def test_update_team_standalone_unchanged_budget_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget", "tpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6810,6 +6839,7 @@ async def test_update_team_standalone_lower_budget_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6912,6 +6942,8 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("max_budget"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6992,6 +7024,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("models"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7091,6 +7124,8 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit( mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("max_budget"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7202,6 +7237,8 @@ async def test_update_team_org_scoped_models_bypasses_user_limit( mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7304,6 +7341,8 @@ async def test_update_team_org_scoped_models_not_in_org_models(): mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7393,6 +7432,8 @@ async def test_update_team_org_scoped_models_with_all_proxy_models( mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7502,6 +7543,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("tpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7584,6 +7626,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("rpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7981,6 +8024,8 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("tpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8067,6 +8112,8 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("rpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8158,6 +8205,8 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit( mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("tpm_limit", "rpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8286,6 +8335,7 @@ async def test_update_team_guardrails_with_org_id( } with ( + _team_admin_may_edit("guardrails", "organization_id"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -11177,8 +11227,8 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) with patch( - "litellm.proxy.management_endpoints.team_endpoints._verify_team_access", - AsyncMock(return_value=None), + "litellm.proxy.management_endpoints.team_endpoints._resolve_team_access", + AsyncMock(return_value="org_admin"), ): with pytest.raises(ProxyException) as exc: await update_team( @@ -13246,6 +13296,7 @@ async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin with contextlib.ExitStack() as stack: _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + stack.enter_context(_team_admin_may_edit("default_estimated_output_tokens")) with pytest.raises(ProxyException) as exc: await update_team( data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1), @@ -13277,6 +13328,7 @@ async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edi with contextlib.ExitStack() as stack: prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + stack.enter_context(_team_admin_may_edit("team_alias")) await update_team( data=UpdateTeamRequest( team_id="test_team_id", @@ -13336,6 +13388,7 @@ async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_a with contextlib.ExitStack() as stack: _wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000}) + stack.enter_context(_team_admin_may_edit("metadata")) with pytest.raises(ProxyException) as exc: await update_team( data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}), @@ -14894,3 +14947,178 @@ async def test_update_team_model_max_budget_raise_blocked_for_team_admin(): assert exc.value.code == "403" assert "proxy admin" in str(exc.value.message).lower() mock_prisma.db.litellm_teamtable.update.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# LIT-5722: team admins reach update_team through self_managed_routes and are +# filtered by the team_admin_editable_team_fields setting. +# --------------------------------------------------------------------------- + +_TEAM_ADMIN_CALLER = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team-admin", user_id="team-admin" +) +_PROXY_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin") + + +def _update_request_stub(): + from unittest.mock import Mock + + from fastapi import Request + + return Mock(spec=Request) + + +@pytest.mark.asyncio +async def test_update_team_team_admin_is_refused_before_any_write_when_no_fields_are_enabled(): + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit()) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "cannot edit team settings" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_configured_but_unsupported_field_does_not_open_editing(): + """Only fields in SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS count, whatever general_settings says.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context( + patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["team_alias"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global + ) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "cannot edit team settings" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_team_admin_changing_an_unpermitted_field_is_refused_by_name(): + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit("team_alias")) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=10), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "'tpm_limit'" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_allowed( + disable_audit_logging_for_mocked_team, +): + """The dashboard resends the whole form, so only a value that differs from what is stored counts.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit("team_alias")) + result = await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=None, models=[]), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert result["data"].team_id == "test_team_id" + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list( + disable_audit_logging_for_mocked_team, +): + """A caller who is both org admin and roster admin keeps unrestricted edits.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit()) + stack.enter_context( + patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=True), + ) + ) + result = await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert result["data"].team_id == "test_team_id" + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_unknown_team_is_403_for_non_proxy_admins_and_404_for_proxy_admins(): + """Now that any authenticated caller reaches the handler, 'team not found' must not leak team ids.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(ProxyException) as denied: + await update_team( + data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + with pytest.raises(ProxyException) as missing: + await update_team( + data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_PROXY_ADMIN_CALLER, + ) + + assert str(denied.value.code) == "403" + assert "do not have access to this team" in str(denied.value.message) + assert "no-such-team" not in str(denied.value.message) + assert str(missing.value.code) == "404" + + +@pytest.mark.asyncio +async def test_resolve_team_access_ranks_proxy_admin_then_org_admin_then_team_admin(): + from litellm.proxy.management_endpoints.team_endpoints import _resolve_team_access + + team = LiteLLM_TeamTable( + team_id="team-1", + organization_id="org-1", + members_with_roles=[Member(user_id="team-admin", role="admin")], + ) + roster_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin") + outsider = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else") + org_lookup = AsyncMock(return_value=False) + + with patch("litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", org_lookup): # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + assert await _resolve_team_access(team_obj=team, user_api_key_dict=_PROXY_ADMIN_CALLER) == "proxy_admin" + assert org_lookup.await_count == 0 + assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "team_admin" + assert await _resolve_team_access(team_obj=team, user_api_key_dict=outsider) is None + org_lookup.return_value = True + assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "org_admin" diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 709447d23c0..e846b8a5111 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3266,3 +3266,111 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 400 assert "enable_ptu_cost_attribution" in str(response.json()["detail"]) assert not mock_prisma.db.litellm_uisettings.upsert.called + + +class TestTeamAdminEditableTeamFieldsSetting: + """team_admin_editable_team_fields: the proxy-wide allow-list update_team applies to team admins.""" + + def _as_proxy_admin(self, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + return mock_prisma + + def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset({"tpm_limit"}), + ) + + try: + response = client.patch( + "/update/ui_settings", + json={"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 400 + detail = response.json()["detail"]["error"] + assert "['blocked', 'organization_id']" in detail + assert "['tpm_limit']" in detail + assert not mock_prisma.db.litellm_uisettings.upsert.called + + def test_patch_rejects_a_non_list_value(self, monkeypatch): + self._as_proxy_admin(monkeypatch) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": "tpm_limit"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 422 + + def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset({"tpm_limit", "rpm_limit"}), + ) + general_settings: dict = {"team_admin_editable_team_fields": ["rpm_limit"]} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["tpm_limit"]}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == ["tpm_limit"] + assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"] + + def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": []}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == [] + assert general_settings["team_admin_editable_team_fields"] == [] + + def test_get_reports_the_stored_list_and_advertises_supported_fields(self, mock_auth, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.ui_settings = {"team_admin_editable_team_fields": ["tpm_limit"]} + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + general_settings: dict = {} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + data = response.json() + assert data["values"]["team_admin_editable_team_fields"] == ["tpm_limit"] + assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"] + field_schema = data["field_schema"]["properties"]["team_admin_editable_team_fields"] + assert field_schema["type"] == "array" + assert field_schema["items"]["type"] == "string" + assert isinstance(field_schema["items"]["enum"], list) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx new file mode 100644 index 00000000000..d1a44460190 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx @@ -0,0 +1,91 @@ +import { describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; + +import { renderWithProviders, screen } from "@/../tests/test-utils"; + +import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; + +describe("TeamAdminEditableFieldsSettings", () => { + it("explains that nothing can be enabled when the proxy supports no fields", () => { + renderWithProviders( + , + ); + + expect(screen.getByText("Team admins cannot edit team settings")).toBeInTheDocument(); + expect(screen.getByText(/does not support enabling any team settings fields/)).toBeInTheDocument(); + expect(screen.queryByRole("checkbox")).not.toBeInTheDocument(); + }); + + it("renders one checkbox per supported field, checked for the enabled ones", () => { + renderWithProviders( + , + ); + + expect(screen.getByText("1 field enabled")).toBeInTheDocument(); + expect(screen.getByText("Fields a team admin may change")).toBeInTheDocument(); + expect(screen.getByRole("checkbox", { name: "max_budget" })).not.toBeChecked(); + expect(screen.getByRole("checkbox", { name: "tpm_limit" })).toBeChecked(); + }); + + it("saves the list with the field added when an unchecked field is ticked", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "max_budget" })); + + expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["tpm_limit", "max_budget"] }); + }); + + it("saves the list with the field removed when a checked field is unticked", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "tpm_limit" })); + + expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }); + }); + + it("blocks toggling while a save is in flight", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "tpm_limit" })); + + expect(onUpdate).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx new file mode 100644 index 00000000000..8a42835ac4e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx @@ -0,0 +1,64 @@ +"use client"; + +import { Badge } from "@/components/ui/badge"; +import { Checkbox } from "@/components/ui/checkbox"; + +interface TeamAdminEditableFieldsSettingsProps { + editableFields: readonly string[]; + supportedFields: readonly string[]; + description?: string; + isUpdating: boolean; + onUpdate: (settings: { team_admin_editable_team_fields: string[] }) => void; +} + +export default function TeamAdminEditableFieldsSettings({ + editableFields, + supportedFields, + description, + isUpdating, + onUpdate, +}: TeamAdminEditableFieldsSettingsProps) { + const toggleField = (field: string, checked: boolean) => { + const next = checked ? [...editableFields, field] : editableFields.filter((item) => item !== field); + onUpdate({ team_admin_editable_team_fields: next }); + }; + + return ( +
+
+
+

Team admin editable fields

+ 0 ? "secondary" : "outline"}> + {editableFields.length > 0 + ? `${editableFields.length} field${editableFields.length !== 1 ? "s" : ""} enabled` + : "Team admins cannot edit team settings"} + +
+ {description &&

{description}

} +
+ + {supportedFields.length === 0 ? ( +

+ This proxy version does not support enabling any team settings fields for team admins yet. +

+ ) : ( +
+ {supportedFields.map((field) => { + const checkboxId = `team-admin-editable-${field}`; + return ( + + ); + })} +
+ )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx index c2834e65498..26d193cf5d8 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx @@ -155,4 +155,48 @@ describe("UISettings", () => { ); expect(toast.success).toHaveBeenCalledWith("UI settings updated successfully"); }); + + it("saves the team admin editable field list when a supported field is ticked", () => { + const mutateMock = vi.fn((_settings, options) => { + options?.onSuccess?.(); + }); + mockUseUpdateUISettings.mockReturnValue({ + mutate: mutateMock, + isPending: false, + error: null, + }); + mockUseUISettings.mockReturnValue( + buildSettingsResponse({ + data: { + field_schema: { + properties: { + team_admin_editable_team_fields: { + description: "Team settings fields a team admin may change", + type: "array", + items: { type: "string", enum: ["tpm_limit"] }, + }, + }, + }, + values: { team_admin_editable_team_fields: [] }, + }, + }), + ); + + render(); + + expect(screen.getByText("Team settings fields a team admin may change")).toBeInTheDocument(); + + act(() => { + fireEvent.click(screen.getByRole("checkbox", { name: "tpm_limit" })); + }); + + expect(mutateMock).toHaveBeenCalledWith( + { team_admin_editable_team_fields: ["tpm_limit"] }, + expect.objectContaining({ + onSuccess: expect.any(Function), + onError: expect.any(Function), + }), + ); + expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully"); + }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 612ca05d083..04c53ec39e8 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -9,7 +9,12 @@ import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Separator } from "@/components/ui/separator"; import { Skeleton } from "@/components/ui/skeleton"; import { Switch } from "@/components/ui/switch"; +import { + parseSupportedTeamAdminEditableFields, + parseTeamAdminEditableFields, +} from "@/components/team/teamAdminEditAccess"; import PageVisibilitySettings from "./PageVisibilitySettings"; +import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; interface SettingRowProps { ariaLabel: string; @@ -65,6 +70,7 @@ export default function UISettings() { const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; + const teamAdminEditableFieldsProperty = schema?.properties?.team_admin_editable_team_fields; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -110,6 +116,17 @@ export default function UISettings() { }); }; + const handleUpdateTeamAdminEditableFields = (settings: { team_admin_editable_team_fields: string[] }) => { + updateSettings(settings, { + onSuccess: () => { + toast.success("Team admin editable fields updated successfully"); + }, + onError: (error) => { + toast.fromError(error); + }, + }); + }; + const handleToggleForwardClientHeaders = (checked: boolean) => { updateSettings( { forward_client_headers_to_llm_api: checked }, @@ -439,6 +456,15 @@ export default function UISettings() { isUpdating={isUpdating} onUpdate={handleUpdatePageVisibility} /> + + + )} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 16eb70fdf5e..00bf264c636 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -69,6 +69,10 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), })); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: vi.fn(), +})); + vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); @@ -228,6 +232,7 @@ import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; const mockUseAllProxyModels = vi.mocked(useAllProxyModels); const mockUseKeys = vi.mocked(useKeys); @@ -237,6 +242,7 @@ const mockUseCurrentUser = vi.mocked(useCurrentUser); const mockUseMCPServers = vi.mocked(useMCPServers); const mockUseMCPToolsets = vi.mocked(useMCPToolsets); const mockUseAccessGroups = vi.mocked(useAccessGroups); +const mockUseUISettings = vi.mocked(useUISettings); const createMockTeamData = (overrides = {}) => ({ team_id: "123", @@ -305,6 +311,10 @@ const seedDefaultMocks = () => { isLoading: false, isError: false, } as any); + mockUseUISettings.mockReturnValue({ + data: { values: { team_admin_editable_team_fields: [] } }, + isLoading: false, + } as any); mockUseKeys.mockReturnValue({ data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, isPending: false, @@ -1863,6 +1873,61 @@ describe("TeamInfoView", () => { }); }); }); + + describe("team admin edit access", () => { + const teamAdminProps = { ...defaultProps, is_proxy_admin: false, is_team_admin: true }; + + beforeEach(() => { + authState.userRole = "Internal User"; + }); + + it("tells a team admin to ask a proxy admin when no team field is enabled for them", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + + expect(toast.error).toHaveBeenCalledWith("Team admins cannot edit team settings on this proxy", { + description: "Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.", + }); + expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + it("opens the form for a team admin once a proxy admin has enabled a field", async () => { + mockUseUISettings.mockReturnValue({ + data: { values: { team_admin_editable_team_fields: ["tpm_limit"] } }, + isLoading: false, + } as any); + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + + expect(await screen.findByLabelText("Team Name")).toBeInTheDocument(); + expect(toast.error).not.toHaveBeenCalled(); + }); + + it("never gates a proxy admin on the team admin field list", async () => { + authState.userRole = "Admin"; + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + + expect(await screen.findByLabelText("Team Name")).toBeInTheDocument(); + expect(toast.error).not.toHaveBeenCalled(); + }); + }); }); describe("TeamInfoView - which team member fields reach the update payload depends on the open sections", () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index c476f3492a3..af94b26ed51 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1,6 +1,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { useQueryClient } from "@tanstack/react-query"; import UserSearchModal from "@/components/common_components/user_search_modal"; import { @@ -48,6 +49,11 @@ import React, { useEffect, useMemo, useState } from "react"; import { useFieldArray } from "react-hook-form"; import { z } from "zod/v4"; import GuardrailsSelect from "./GuardrailsSelect"; +import { + resolveTeamEditAccess, + TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION, + TEAM_ADMIN_EDITING_DISABLED_TITLE, +} from "./teamAdminEditAccess"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; @@ -582,6 +588,7 @@ const TeamInfoView: React.FC = ({ const canEditTeamEstimates = isProxyAdminRole(userRole); const teamEstimateTooltip = estimateTooltips(canEditTeamEstimates, "team"); const { data: userOrganizations = [] } = useOrganizations(); + const { data: uiSettingsData } = useUISettings(); const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema(); const queryClient = useQueryClient(); @@ -625,6 +632,12 @@ const TeamInfoView: React.FC = ({ ); const canEditTeam = is_team_admin || is_proxy_admin || is_org_admin || isOrgAdminForTeam || isTeamAdminFromTeamData; + const editsAsTeamAdmin = + canEditTeam && !is_proxy_admin && !isProxyAdminRole(userRole) && !is_org_admin && !isOrgAdminForTeam; + const teamEditAccess = useMemo( + () => resolveTeamEditAccess(editsAsTeamAdmin, uiSettingsData?.values), + [editsAsTeamAdmin, uiSettingsData], + ); const visibleTabs = useMemo(() => getTeamInfoVisibleTabs(canEditTeam), [canEditTeam]); const defaultTabKey = useMemo(() => getTeamInfoDefaultTab(editTeam, canEditTeam), [editTeam, canEditTeam]); const { onTabChange, hasVisited } = useVisitedTabs(defaultTabKey); @@ -644,6 +657,15 @@ const TeamInfoView: React.FC = ({ setIsEditing(true); }; + const openSettingsEditor = (modelAliases: Record) => { + if (teamEditAccess.kind === "team_admin_disabled") { + toast.error(TEAM_ADMIN_EDITING_DISABLED_TITLE, { description: TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION }); + return; + } + setTeamModelAliases(modelAliases); + startEditing(); + }; + const applyKillSwitchToGuardrails = (checked: boolean) => { const current = form.getValues("guardrails") ?? []; const nonGlobals = current.filter((name) => !globalGuardrailNames.has(name)); @@ -1340,10 +1362,7 @@ const TeamInfoView: React.FC = ({ {canEditTeam && !isEditing && ( + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index dea53eb42a2..e954bc1c581 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1889,18 +1889,32 @@ describe("TeamInfoView", () => { expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); }); - it("opens the form for a team admin once the proxy reports an enabled field", async () => { + it("gives a team admin only the fields the proxy enabled and sends only those on save", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ caller_edit_access: { kind: "team_admin", editable_fields: ["tpm_limit"] } }), + createMockTeamData({ + tpm_limit: 1000, + caller_edit_access: { kind: "team_admin", editable_fields: ["tpm_limit"] }, + }), ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); renderWithProviders(); await user.click(await screen.findByRole("tab", { name: "Settings" })); await user.click(await screen.findByRole("button", { name: /edit settings/i })); - expect(await screen.findByLabelText("Team Name")).toBeInTheDocument(); + const tpmInput = await screen.findByLabelText("Tokens per minute Limit (TPM)"); + expect(tpmInput).toHaveValue(1000); + expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Requests per minute Limit (RPM)")).not.toBeInTheDocument(); + + fireEvent.change(tpmInput, { target: { value: "5000" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(1)); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).toStrictEqual({ team_id: "123", tpm_limit: 5000 }); + expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully"); expect(toast.error).not.toHaveBeenCalled(); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 48de90f74b5..30b648fc53c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -53,7 +53,9 @@ import { parseTeamEditAccess, TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION, TEAM_ADMIN_EDITING_DISABLED_TITLE, + type TeamAdminSettingsChanges, } from "./teamAdminEditAccess"; +import TeamAdminSettingsForm from "./TeamAdminSettingsForm"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; @@ -862,6 +864,27 @@ const TeamInfoView: React.FC = ({ setMemberToDelete(null); }; + const persistTeamUpdate = async (token: string, updateData: Record) => { + await teamUpdateCall(token, updateData); + queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + + toast.success("Team settings updated successfully"); + setIsEditing(false); + fetchTeamInfo(); + }; + + const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => { + if (!accessToken) return; + setIsTeamSaving(true); + try { + await persistTeamUpdate(accessToken, { team_id: teamId, ...changes }); + } catch (error) { + console.error("Error updating team:", error); + } finally { + setIsTeamSaving(false); + } + }; + const handleTeamUpdate = async (values: any) => { try { if (!accessToken) return; @@ -1112,12 +1135,7 @@ const TeamInfoView: React.FC = ({ } } - await teamUpdateCall(accessToken, updateData); - queryClient.invalidateQueries({ queryKey: organizationKeys.all }); - - toast.success("Team settings updated successfully"); - setIsEditing(false); - fetchTeamInfo(); + await persistTeamUpdate(accessToken, updateData); } catch (error) { console.error("Error updating team:", error); } finally { @@ -1135,6 +1153,17 @@ const TeamInfoView: React.FC = ({ const { team_info: info } = teamData; + const teamAdminSettingsEditor = + teamEditAccess.kind === "team_admin" ? ( + setIsEditing(false)} + onSave={saveTeamAdminSettings} + /> + ) : null; + const inheritedMcpServers = computeInheritedGrants( info.access_group_mcp_server_ids, info.access_group_details, @@ -1347,8 +1376,8 @@ const TeamInfoView: React.FC = ({ )} - {isEditing && isGuardrailsLoading ? ( -
Loading...
+ {isEditing && (teamAdminSettingsEditor !== null || isGuardrailsLoading) ? ( + teamAdminSettingsEditor ??
Loading...
) : isEditing ? (
void form.handleSubmit(onTeamUpdateSubmit)(event)}> diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index 122cc749e6b..c8117800053 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -4,8 +4,40 @@ import { parseSupportedTeamAdminEditableFields, parseTeamAdminEditableFields, parseTeamEditAccess, + teamAdminFieldLabel, + teamAdminSettingsChanges, } from "./teamAdminEditAccess"; +describe("teamAdminFieldLabel", () => { + it("names tpm_limit the way the team settings form does", () => { + expect(teamAdminFieldLabel("tpm_limit")).toBe("Tokens per minute Limit (TPM)"); + }); + + it("falls back to the raw field name for a field the dashboard has no label for", () => { + expect(teamAdminFieldLabel("max_budget")).toBe("max_budget"); + }); +}); + +describe("teamAdminSettingsChanges", () => { + const tpmEnabled = new Set(["tpm_limit"]); + + it.each([ + ["a typed number string", "5000", 5000], + ["a stored number", 1200, 1200], + ["zero", "0", 0], + ["an emptied input", "", null], + ["whitespace", " ", null], + ["no stored limit", null, null], + ["an unset value", undefined, null], + ])("sends tpm_limit for %s", (_label, tpm_limit, expected) => { + expect(teamAdminSettingsChanges({ tpm_limit }, tpmEnabled)).toStrictEqual({ tpm_limit: expected }); + }); + + it("leaves tpm_limit out when the proxy did not enable it for team admins", () => { + expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, new Set(["max_budget"]))).toStrictEqual({}); + }); +}); + describe("parseTeamAdminEditableFields", () => { it("returns the configured list", () => { expect(parseTeamAdminEditableFields({ team_admin_editable_team_fields: ["tpm_limit", "rpm_limit"] })).toEqual([ diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index 79cc81b1416..d38566eefda 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -39,6 +39,29 @@ export const parseSupportedTeamAdminEditableFields = (uiSettingsFieldSchema: unk return items.success ? fieldListSchema.parse(items.data.enum) : []; }; +const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([["tpm_limit", "Tokens per minute Limit (TPM)"]]); + +export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field; + +export interface TeamAdminSettingsValues { + readonly tpm_limit?: string | number | null; +} + +export interface TeamAdminSettingsChanges { + readonly tpm_limit?: number | null; +} + +const numberOrNull = (value: string | number | null | undefined): number | null => { + if (value === null || value === undefined || String(value).trim() === "") return null; + const parsed = Number(value); + return Number.isNaN(parsed) ? null : parsed; +}; + +export const teamAdminSettingsChanges = ( + values: TeamAdminSettingsValues, + editableFields: ReadonlySet, +): TeamAdminSettingsChanges => (editableFields.has("tpm_limit") ? { tpm_limit: numberOrNull(values.tpm_limit) } : {}); + export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => { const parsed = callerEditAccessSchema.safeParse(callerEditAccess); if (!parsed.success) return { kind: "none" }; From 6e2ae196705e6bc9745488238e47623ec55d2f60 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 11:09:44 -0700 Subject: [PATCH 63/89] fix(proxy): enforce org budget ceilings on /team/update update_team loaded the org without its budget row, so the org max_budget, tpm_limit and rpm_limit checks silently passed. It now loads the budget the same way /team/new does --- .../management_endpoints/team_endpoints.py | 1 + .../management/test_team_budget_limits.py | 45 ++++++++++--- .../test_team_endpoints.py | 63 +++++++++++++++++++ 3 files changed, 100 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6fccdf02123..f0ec3f975d5 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2330,6 +2330,7 @@ async def update_team( org_id=org_id_to_check, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, + include_budget_table=True, ) if org_table is not None: await _check_org_team_limits( diff --git a/tests/proxy_behavior/management/test_team_budget_limits.py b/tests/proxy_behavior/management/test_team_budget_limits.py index e5c67ee4e39..a172f625e91 100644 --- a/tests/proxy_behavior/management/test_team_budget_limits.py +++ b/tests/proxy_behavior/management/test_team_budget_limits.py @@ -10,12 +10,11 @@ Pins the five helpers Driven through /team/new + /team/update. -Structural finding, updated: /team/new loads the org via `get_org_object` -WITH `include_budget_table=True`, so the org max_budget / org tpm / org rpm -guards inside `_check_org_team_limits` are live there and are pinned as -enforced below. /team/update still loads the org without the budget -relation, so its budget guards remain no-ops. The `models` subset guard IS -reachable on both because it reads `org_table.models` directly. The +Structural finding, updated: /team/new and /team/update both load the org +via `get_org_object` WITH `include_budget_table=True`, so the org max_budget / +org tpm / org rpm guards inside `_check_org_team_limits` are live on both and +are pinned as enforced below. The `models` subset guard reads +`org_table.models` directly. The `_check_user_team_limits` guards reach all branches through `user_api_key_dict`, no relation include needed. """ @@ -139,9 +138,8 @@ async def test_check_org_team_limits_models_subset( # --------------------------------------------------------------------------- -# _check_org_team_limits — budget / tpm / rpm live on /team/new since its -# get_org_object call passes include_budget_table=True. (/team/update still -# loads the org without the budget relation, so its guards remain no-ops.) +# _check_org_team_limits — budget / tpm / rpm live on /team/new and +# /team/update since both get_org_object calls pass include_budget_table=True. # --------------------------------------------------------------------------- _ORG_BUDGET_ENFORCED_SCENARIOS = [ @@ -216,6 +214,35 @@ async def test_check_org_team_limits_budget_enforced( assert len(rows) == (1 if expected_status == 200 else 0) +@pytest.mark.parametrize( + "org_budget,body_extras,expected_status", + [(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS], + ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS], +) +async def test_check_org_team_limits_budget_enforced_on_update( + org_budget, + body_extras: Dict[str, Any], + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget) + team_id = await create_scratch_team(prisma, scratch.tag("team"), organization_id=org_id) + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {seeder}"}, + json={"team_id": team_id, **body_extras}, + ) + assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}" + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + persisted = {field: getattr(row, field) for field in body_extras} + assert (persisted == body_extras) == (expected_status == 200) + + # --------------------------------------------------------------------------- # _check_user_team_limits — fires for standalone (no-org) teams created by # a non-admin caller. Each guard reads from user_api_key_dict / user_obj. 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 6cece466943..f507311a24f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15078,6 +15078,69 @@ async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enabl assert "'rpm_limit'" in str(refused.value.message) +@pytest.mark.asyncio +async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit_logging_for_mocked_team): + """The org ceiling lives on the org's budget row, so /team/update must load it to enforce the cap.""" + import contextlib + + capped_org = LiteLLM_OrganizationTable( + organization_id="capped-org", + budget_id="capped-budget", + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=10000), + ) + + async def org_lookup(**kwargs): + return capped_org if kwargs.get("include_budget_table") else capped_org.model_copy( + update={"litellm_budget_table": None} + ) + + org_team = MagicMock() + org_team.metadata = {} + org_team.organization_id = "capped-org" + org_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "organization_id": "capped-org", + "metadata": {}, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=org_team) + stack.enter_context(_team_admin_may_edit("tpm_limit")) + stack.enter_context( + patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=False), + ) + ) + stack.enter_context( + patch( # test-quality-ok: update_team reads orgs through this module-level import; no seam to inject + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + AsyncMock(side_effect=org_lookup), + ) + ) + with pytest.raises(ProxyException) as over_cap: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=20000), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=8000), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(over_cap.value.code) == "400" + assert "exceeds organization's tpm_limit (10000)" in str(over_cap.value.message) + assert prisma.db.litellm_teamtable.update.await_count == 1 + assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 8000 + + @pytest.mark.asyncio async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list( disable_audit_logging_for_mocked_team, From b92bd98df6f234f3acc06a9da6d62a848bbbf08e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:31:02 -0700 Subject: [PATCH 64/89] feat(ui): save team admin editable fields from their own card with a Save button Ticking a field only changes the draft. The allow-list is written when the proxy admin clicks Save, and the card sits next to UI Settings instead of inside its auto-saving toggles. --- .../admin-panel/_components/AdminPanel.tsx | 2 + .../TeamAdminEditableFieldsSettings.test.tsx | 208 ++++++++++++------ .../TeamAdminEditableFieldsSettings.tsx | 173 ++++++++++----- .../UISettings/UISettings.test.tsx | 44 ---- .../AdminSettings/UISettings/UISettings.tsx | 26 --- 5 files changed, 270 insertions(+), 183 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index 1f35f46dcd4..1c8425251fd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -17,6 +17,7 @@ import SCIMConfig from "@/components/SCIM"; import LoggingSettings from "@/components/Settings/AdminSettings/LoggingSettings/LoggingSettings"; import SSOSettings from "@/components/Settings/AdminSettings/SSOSettings/SSOSettings"; import UISettings from "@/components/Settings/AdminSettings/UISettings/UISettings"; +import TeamAdminEditableFieldsSettings from "@/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings"; import UserBannerSettings from "@/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings"; import CyberArk from "@/components/Settings/AdminSettings/CyberArk/CyberArk"; import HashicorpVault from "@/components/Settings/AdminSettings/HashicorpVault/HashicorpVault"; @@ -382,6 +383,7 @@ const AdminPanel: React.FC = ({ proxySettings }) => { children: (
+
), diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx index 3880acfb6d5..2e1a8e9fd36 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx @@ -1,91 +1,175 @@ -import { describe, expect, it, vi } from "vitest"; -import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen } from "@/../tests/test-utils"; +import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; +import { toast } from "@/lib/toast"; import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; +const mockUseUISettings = vi.hoisted(() => vi.fn()); +const mockUseUpdateUISettings = vi.hoisted(() => vi.fn()); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "test-token" }), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: mockUseUISettings, +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings", () => ({ + useUpdateUISettings: mockUseUpdateUISettings, +})); + +const TPM_LABEL = "Tokens per minute Limit (TPM)"; + +const mockSettings = (supported: readonly string[], enabled: readonly string[]) => + mockUseUISettings.mockReturnValue({ + isLoading: false, + data: { + field_schema: { + properties: { + team_admin_editable_team_fields: { + description: "Fields a team admin may change", + items: { type: "string", enum: supported }, + }, + }, + }, + values: { team_admin_editable_team_fields: enabled }, + }, + }); + +const mockSave = ({ + isPending = false, + outcome = "success", +}: { + isPending?: boolean; + outcome?: "success" | "error"; +}) => { + const mutate = vi.fn((_settings: unknown, options: { onSuccess: () => void; onError: (error: Error) => void }) => + outcome === "success" ? options.onSuccess() : options.onError(new Error("save failed")), + ); + mockUseUpdateUISettings.mockReturnValue({ mutate, isPending }); + return mutate; +}; + +const saveButton = () => screen.getByRole("button", { name: "Save" }); + describe("TeamAdminEditableFieldsSettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + it("explains that nothing can be enabled when the proxy supports no fields", () => { - renderWithProviders( - , - ); + mockSettings([], []); + mockSave({}); + + renderWithProviders(); expect(screen.getByText("Team admins cannot edit team settings")).toBeInTheDocument(); expect(screen.getByText(/does not support enabling any team settings fields/)).toBeInTheDocument(); expect(screen.queryByRole("checkbox")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save" })).not.toBeInTheDocument(); }); - it("renders one checkbox per supported field, checked for the enabled ones and named by the field's form label when the dashboard has one", () => { - renderWithProviders( - , - ); + it("renders one checkbox per supported field, checked for the saved ones, with Save disabled until something changes", () => { + mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]); + mockSave({}); + renderWithProviders(); + + expect(screen.getByText("Team admin editable fields")).toBeInTheDocument(); expect(screen.getByText("1 field enabled")).toBeInTheDocument(); expect(screen.getByText("Fields a team admin may change")).toBeInTheDocument(); expect(screen.getByRole("checkbox", { name: "max_budget" })).not.toBeChecked(); - expect(screen.getByRole("checkbox", { name: "Tokens per minute Limit (TPM)" })).toBeChecked(); + expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked(); + expect(saveButton()).toBeDisabled(); }); - it("saves the list with the field added when an unchecked field is ticked", async () => { - const onUpdate = vi.fn(); - const user = userEvent.setup(); - renderWithProviders( - , + it("only saves a ticked field once Save is clicked", async () => { + mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]); + const mutate = mockSave({}); + + renderWithProviders(); + fireEvent.click(screen.getByRole("checkbox", { name: "max_budget" })); + + expect(screen.getByRole("checkbox", { name: "max_budget" })).toBeChecked(); + expect(mutate).not.toHaveBeenCalled(); + + fireEvent.click(saveButton()); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully")); + expect(mutate).toHaveBeenCalledWith( + { team_admin_editable_team_fields: ["max_budget", "tpm_limit"] }, + expect.anything(), ); - - await user.click(screen.getByRole("checkbox", { name: "max_budget" })); - - expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["tpm_limit", "max_budget"] }); + expect(saveButton()).toBeDisabled(); }); - it("saves the list with the field removed when a checked field is unticked", async () => { - const onUpdate = vi.fn(); - const user = userEvent.setup(); - renderWithProviders( - , - ); + it("saves the list without an unticked field", async () => { + mockSettings(["max_budget", "tpm_limit"], ["max_budget", "tpm_limit"]); + const mutate = mockSave({}); - await user.click(screen.getByRole("checkbox", { name: "Tokens per minute Limit (TPM)" })); + renderWithProviders(); + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + fireEvent.click(saveButton()); - expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }); + await waitFor(() => expect(mutate).toHaveBeenCalledTimes(1)); + expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }, expect.anything()); }); - it("blocks toggling while a save is in flight", async () => { - const onUpdate = vi.fn(); - const user = userEvent.setup(); - renderWithProviders( - , - ); + it("disables Save again when the draft is ticked back to the saved list", () => { + mockSettings(["tpm_limit"], []); + mockSave({}); - await user.click(screen.getByRole("checkbox", { name: "Tokens per minute Limit (TPM)" })); + renderWithProviders(); + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); - expect(onUpdate).not.toHaveBeenCalled(); + expect(saveButton()).toBeEnabled(); + + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + + expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked(); + expect(saveButton()).toBeDisabled(); + }); + + it("treats a saved list in another order, or with fields this proxy dropped, as the same selection", () => { + mockSettings(["max_budget", "tpm_limit"], ["tpm_limit", "retired_field", "max_budget"]); + mockSave({}); + + renderWithProviders(); + + expect(screen.getByText("2 fields enabled")).toBeInTheDocument(); + + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + + expect(saveButton()).toBeDisabled(); + }); + + it("keeps the draft and shows the error when the save fails", async () => { + mockSettings(["tpm_limit"], []); + const mutate = mockSave({ outcome: "error" }); + + renderWithProviders(); + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + fireEvent.click(saveButton()); + + await waitFor(() => expect(toast.fromError).toHaveBeenCalledTimes(1)); + expect(mutate).toHaveBeenCalledTimes(1); + expect(toast.success).not.toHaveBeenCalled(); + expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked(); + expect(saveButton()).toBeEnabled(); + }); + + it("blocks ticking and saving while a save is in flight", () => { + mockSettings(["tpm_limit"], []); + const mutate = mockSave({ isPending: true }); + + renderWithProviders(); + fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL })); + + expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked(); + expect(screen.getByRole("button", { name: "Saving..." })).toBeDisabled(); + expect(mutate).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx index 27f53e5fa77..3737d30945e 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx @@ -1,65 +1,136 @@ "use client"; +import { Controller } from "react-hook-form"; +import { z } from "zod/v4"; + +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import { useUpdateUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { + parseSupportedTeamAdminEditableFields, + parseTeamAdminEditableFields, + teamAdminFieldLabel, +} from "@/components/team/teamAdminEditAccess"; import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Checkbox } from "@/components/ui/checkbox"; -import { teamAdminFieldLabel } from "@/components/team/teamAdminEditAccess"; +import { Skeleton } from "@/components/ui/skeleton"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { toast } from "@/lib/toast"; -interface TeamAdminEditableFieldsSettingsProps { - editableFields: readonly string[]; - supportedFields: readonly string[]; - description?: string; - isUpdating: boolean; - onUpdate: (settings: { team_admin_editable_team_fields: string[] }) => void; -} +const editableFieldsSchema = z.object({ team_admin_editable_team_fields: z.array(z.string()) }); -export default function TeamAdminEditableFieldsSettings({ - editableFields, - supportedFields, - description, - isUpdating, - onUpdate, -}: TeamAdminEditableFieldsSettingsProps) { - const toggleField = (field: string, checked: boolean) => { - const next = checked ? [...editableFields, field] : editableFields.filter((item) => item !== field); - onUpdate({ team_admin_editable_team_fields: next }); - }; +type SaveEditableFields = ReturnType["mutate"]; + +export default function TeamAdminEditableFieldsSettings() { + const { accessToken } = useAuthorized(); + const { data, isLoading } = useUISettings(); + const { mutate: saveSettings, isPending } = useUpdateUISettings(accessToken); + const supportedFields = parseSupportedTeamAdminEditableFields(data?.field_schema); + const savedFields = parseTeamAdminEditableFields(data?.values); + const enabledFields = supportedFields.filter((field) => savedFields.includes(field)); return ( -
-
+ +
-

Team admin editable fields

- 0 ? "secondary" : "outline"}> - {editableFields.length > 0 - ? `${editableFields.length} field${editableFields.length !== 1 ? "s" : ""} enabled` + Team admin editable fields + 0 ? "secondary" : "outline"}> + {enabledFields.length > 0 + ? `${enabledFields.length} field${enabledFields.length !== 1 ? "s" : ""} enabled` : "Team admins cannot edit team settings"}
- {description &&

{description}

} -
- - {supportedFields.length === 0 ? ( -

- This proxy version does not support enabling any team settings fields for team admins yet. -

- ) : ( -
- {supportedFields.map((field) => { - const checkboxId = `team-admin-editable-${field}`; - return ( - - ); - })} -
- )} -
+ + {data?.field_schema?.properties?.team_admin_editable_team_fields?.description ?? + "Team settings fields a team admin may change on the teams they administer."} + + + + {isLoading ? ( + + ) : ( + + )} + + + ); +} + +interface TeamAdminEditableFieldsFormProps { + enabledFields: readonly string[]; + supportedFields: readonly string[]; + isPending: boolean; + saveSettings: SaveEditableFields; +} + +function TeamAdminEditableFieldsForm({ + enabledFields, + supportedFields, + isPending, + saveSettings, +}: TeamAdminEditableFieldsFormProps) { + const form = useZodForm(editableFieldsSchema, { + defaultValues: { team_admin_editable_team_fields: [...enabledFields] }, + }); + const submit = form.handleSubmit((values) => + saveSettings(values, { + onSuccess: () => { + form.reset(values); + toast.success("Team admin editable fields updated successfully"); + }, + onError: (error) => { + toast.fromError(error); + }, + }), + ); + + if (supportedFields.length === 0) { + return ( +

+ This proxy version does not support enabling any team settings fields for team admins yet. +

+ ); + } + + return ( + void submit(event)} className="space-y-4"> + ( +
+ {supportedFields.map((name) => { + const checkboxId = `team-admin-editable-${name}`; + return ( + + ); + })} +
+ )} + /> + + ); } diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx index f38caf0df81..c2834e65498 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx @@ -155,48 +155,4 @@ describe("UISettings", () => { ); expect(toast.success).toHaveBeenCalledWith("UI settings updated successfully"); }); - - it("saves the team admin editable field list when a supported field is ticked", () => { - const mutateMock = vi.fn((_settings, options) => { - options?.onSuccess?.(); - }); - mockUseUpdateUISettings.mockReturnValue({ - mutate: mutateMock, - isPending: false, - error: null, - }); - mockUseUISettings.mockReturnValue( - buildSettingsResponse({ - data: { - field_schema: { - properties: { - team_admin_editable_team_fields: { - description: "Team settings fields a team admin may change", - type: "array", - items: { type: "string", enum: ["tpm_limit"] }, - }, - }, - }, - values: { team_admin_editable_team_fields: [] }, - }, - }), - ); - - render(); - - expect(screen.getByText("Team settings fields a team admin may change")).toBeInTheDocument(); - - act(() => { - fireEvent.click(screen.getByRole("checkbox", { name: "Tokens per minute Limit (TPM)" })); - }); - - expect(mutateMock).toHaveBeenCalledWith( - { team_admin_editable_team_fields: ["tpm_limit"] }, - expect.objectContaining({ - onSuccess: expect.any(Function), - onError: expect.any(Function), - }), - ); - expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully"); - }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 04c53ec39e8..612ca05d083 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -9,12 +9,7 @@ import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Separator } from "@/components/ui/separator"; import { Skeleton } from "@/components/ui/skeleton"; import { Switch } from "@/components/ui/switch"; -import { - parseSupportedTeamAdminEditableFields, - parseTeamAdminEditableFields, -} from "@/components/team/teamAdminEditAccess"; import PageVisibilitySettings from "./PageVisibilitySettings"; -import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; interface SettingRowProps { ariaLabel: string; @@ -70,7 +65,6 @@ export default function UISettings() { const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; - const teamAdminEditableFieldsProperty = schema?.properties?.team_admin_editable_team_fields; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -116,17 +110,6 @@ export default function UISettings() { }); }; - const handleUpdateTeamAdminEditableFields = (settings: { team_admin_editable_team_fields: string[] }) => { - updateSettings(settings, { - onSuccess: () => { - toast.success("Team admin editable fields updated successfully"); - }, - onError: (error) => { - toast.fromError(error); - }, - }); - }; - const handleToggleForwardClientHeaders = (checked: boolean) => { updateSettings( { forward_client_headers_to_llm_api: checked }, @@ -456,15 +439,6 @@ export default function UISettings() { isUpdating={isUpdating} onUpdate={handleUpdatePageVisibility} /> - - - )} From d9de2fc5f7cf23f78fed523f7704420a0d0cf680 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:38:41 -0700 Subject: [PATCH 65/89] style(ui): right-align the team admin editable fields Save button --- .../UISettings/TeamAdminEditableFieldsSettings.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx index 3737d30945e..e671a4a47b5 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx @@ -128,9 +128,11 @@ function TeamAdminEditableFieldsForm({ )} /> - +
+ +
); } From cac521246a900637e43222e8d9db13886fd50eaf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:58:22 -0700 Subject: [PATCH 66/89] style(ui): right-align the Save banner button --- .../AdminSettings/UserBannerSettings/UserBannerSettings.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx index 82c25150ec5..21e9ab0f7d9 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx @@ -145,7 +145,7 @@ function UserBannerSettingsForm({ persisted, isLoading, isPending, saveBanner }: )} -
+
From ebae692a0ddc854bebbec41983ad1fe5ba070b66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:45:50 -0700 Subject: [PATCH 67/89] refactor(responses): drop the api_base cast and mark the header merge mutable-ok --- litellm/llms/azure_ai/responses/transformation.py | 6 +++++- litellm/responses/main.py | 9 +++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/responses/transformation.py b/litellm/llms/azure_ai/responses/transformation.py index b61c856f733..66a284c821d 100644 --- a/litellm/llms/azure_ai/responses/transformation.py +++ b/litellm/llms/azure_ai/responses/transformation.py @@ -34,7 +34,11 @@ class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): litellm_params=params.model_dump(), api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)), ) - return {**headers, **auth_headers, "Content-Type": "application/json"} + return { # mutable-ok: the handler updates the returned headers in place per the dict contract + **headers, + **auth_headers, + "Content-Type": "application/json", + } def supports_native_websocket(self) -> bool: return False diff --git a/litellm/responses/main.py b/litellm/responses/main.py index a138d6f8eb3..63bee9f6d99 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -492,6 +492,11 @@ def _resolve_responses_api_provider_config( return OpenAILikeResponsesConfig() +def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None: + api_base: Final = kwargs.get("api_base") + return api_base if isinstance(api_base, str) else None + + def _will_bridge_to_chat_completions( model: str, custom_llm_provider: str | None, @@ -622,7 +627,7 @@ async def aresponses( custom_llm_provider, bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), - cast(str | None, kwargs.get("api_base")), + _api_base_kwarg(kwargs), ), ): ( @@ -792,7 +797,7 @@ def _apply_prompt_management_to_responses_call( custom_llm_provider, use_chat_completions_api, kwargs.get("model_info"), - cast(str | None, kwargs.get("api_base")), + _api_base_kwarg(kwargs), ), ): ( From 4b70696afafcc11793cba55423fa6e77b0a90b29 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:49:25 +0000 Subject: [PATCH 68/89] fix(streaming): estimate interrupted Anthropic stream usage from reasoning_content Interrupted Anthropic streams that die before message_delta were billed at the message_start placeholder (any value above 1 was trusted) or at 0 when the partial response was reasoning-only, because the token_counter fallback only looked at visible text. Reset the placeholder whenever no finish_reason or second usage event arrived, fold the already-counted reasoning tokens into the fallback estimate, and drop the stale completion_tokens_details so cost is computed from the recovered count Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 70 +++++++++++---- .../test_streaming_chunk_builder_cursor.py | 86 +++++++++++++++++-- 2 files changed, 132 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..a8b1f81702c 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -5,6 +5,7 @@ from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast +from pydantic import BaseModel from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger @@ -148,6 +149,7 @@ class _ToolCallChunk(TypedDict): class _UsageBearingChunk(TypedDict, total=False): usage: Usage | None _hidden_params: Mapping[str, str] + choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]] class _UsageSummary(TypedDict): @@ -921,21 +923,22 @@ class ChunkProcessor: prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details) - completion_tokens = self._reset_anthropic_cursor_completion_tokens( + recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens( chunks=chunks, completion_tokens=completion_tokens, completion_usage_updates=completion_usage_updates, ) + cursor_was_reset: Final = recovered_completion_tokens != completion_tokens return UsagePerChunk( prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + completion_tokens=recovered_completion_tokens, cache_creation_input_tokens=cache_creation_input_tokens, cache_read_input_tokens=cache_read_input_tokens, server_tool_use=server_tool_use, web_search_requests=web_search_requests, google_maps_grounding_requests=google_maps_grounding_requests, - completion_tokens_details=completion_tokens_details, + completion_tokens_details=None if cursor_was_reset else completion_tokens_details, prompt_tokens_details=prompt_tokens_details, cost=cost, inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"), @@ -960,6 +963,32 @@ class ChunkProcessor: ] return values[-1] if values else None + @staticmethod + def _finish_reason_of_choice(choice: object) -> str | None: + match choice: + case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason): + return reason + case {"finish_reason": str() as reason}: + return reason + case _: + return None + + @staticmethod + def _chunk_choices(chunk: "_UsageBearingChunk | BaseModel") -> Sequence[object]: + if isinstance(chunk, dict): + return chunk.get("choices", ()) + if isinstance(chunk, (ModelResponse, ModelResponseStream)): + return chunk.choices + return () + + @staticmethod + def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool: + return any( + ChunkProcessor._finish_reason_of_choice(choice) is not None + for chunk in chunks + for choice in ChunkProcessor._chunk_choices(chunk) + ) + @staticmethod def _reset_anthropic_cursor_completion_tokens( chunks: Sequence["_UsageBearingChunk | ModelResponse"], @@ -970,18 +999,18 @@ class ChunkProcessor: See the ``completion_usage_updates`` comment in ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale - cursor when either it is > 1 (definitely not a placeholder) or we saw - >= 2 completion-bearing usage events (positive evidence ``message_delta`` - arrived). Otherwise — the only completion update we ever saw was the - Anthropic ``message_start`` cursor (=1) — reset to 0 so - ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates - from the actually-received completion text instead of trusting the - placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the - heuristic (which encodes Anthropic's specific message_start SSE shape) - does not silently affect other providers that may legitimately report - ``completion_tokens=1`` from a single usage event. + cursor when we saw >= 2 completion-bearing usage events or any chunk + carried a ``finish_reason`` (positive evidence ``message_delta`` + arrived). Otherwise the only completion update we ever saw was the + Anthropic ``message_start`` cursor, a small placeholder whose magnitude + varies per request (1 and 8 both observed live), so reset to 0 and let + ``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from + the actually-received text and reasoning instead. Gated on + ``custom_llm_provider == "anthropic"`` so the heuristic (which encodes + Anthropic's specific message_start SSE shape) does not silently affect + other providers that legitimately report usage from a single event. """ - saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2 + saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks) if saw_non_cursor_completion: return completion_tokens @@ -995,7 +1024,7 @@ class ChunkProcessor: if isinstance(hp, dict): custom_llm_provider = hp.get("custom_llm_provider") - if custom_llm_provider == "anthropic" and completion_tokens == 1: + if custom_llm_provider == "anthropic": return 0 return completion_tokens @@ -1039,10 +1068,13 @@ class ChunkProcessor: returned_usage.prompt_tokens = 0 returned_usage.completion_tokens = ( completion_tokens - or token_counter( - model=model, - text=completion_output, - count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + or ( + token_counter( + model=model, + text=completion_output, + count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + ) + + (reasoning_tokens or 0) ) ) returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index 3d9971034ae..f4dbb28533f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -22,9 +22,10 @@ text-based fallback to estimate from the real completion text. import pytest - +import litellm from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor from litellm.types.utils import ( + CompletionTokensDetailsWrapper, Delta, ModelResponseStream, StreamingChoices, @@ -35,6 +36,7 @@ from litellm.types.utils import ( def _make_chunk( *, content: str = "", + reasoning_content: str | None = None, usage: Usage = None, finish_reason: str = None, custom_llm_provider: str = "anthropic", @@ -48,7 +50,7 @@ def _make_chunk( StreamingChoices( finish_reason=finish_reason, index=0, - delta=Delta(content=content, role="assistant"), + delta=Delta(content=content, role="assistant", reasoning_content=reasoning_content), ) ], usage=usage, @@ -253,6 +255,79 @@ class TestAnthropicCursorBug: "Reset to 0 forces token_counter fallback." ) + @pytest.mark.parametrize("placeholder", [1, 3, 8]) + def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int): + """ + message_start placeholders are not always 1 (live Anthropic streams + have been observed sending 1 and 8 for the same prompt), and a thinking + model cut off before message_delta has streamed only reasoning_content. + The recovered usage, including the completion_tokens_details the cost + calculator bills from, must come from that reasoning rather than from + the placeholder. + """ + message_start = _make_chunk( + usage=Usage( + prompt_tokens=100, + completion_tokens=placeholder, + total_tokens=100 + placeholder, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=placeholder), + ) + ) + reasoning_text = "Let me work through the scheduling constraints step by step. " * 40 + reasoning_chunks = [ + _make_chunk(reasoning_content=reasoning_text[i : i + 50]) for i in range(0, len(reasoning_text), 50) + ] + + response = litellm.stream_chunk_builder( + chunks=[message_start, *reasoning_chunks], + messages=[{"role": "user", "content": "Plan the schedule."}], + ) + + assert response.choices[0].message.reasoning_content == reasoning_text + reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens + assert reasoning_tokens > placeholder + assert response.usage.completion_tokens == reasoning_tokens, ( + f"Expected completion_tokens to be the reasoning estimate, got " + f"completion_tokens={response.usage.completion_tokens} reasoning_tokens={reasoning_tokens}" + ) + assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens + details = response.usage.completion_tokens_details + assert (details.text_tokens or 0) + details.reasoning_tokens == response.usage.completion_tokens + + def test_fallback_counts_reasoning_and_text_together(self): + """ + With no usable provider count, the estimate covers everything the + provider generated: reasoning_content plus visible text, not text alone. + """ + reasoning = "First I should check whether the input is sorted. " * 10 + text = "The list is already sorted, so no work is needed." + chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)] + + response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}]) + + text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True) + reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens + assert reasoning_tokens > 0 + assert response.usage.completion_tokens == text_only + reasoning_tokens + + def test_lone_usage_event_with_finish_reason_is_trusted(self): + """ + Guardrails rebuild responses from the chunks yielded to the client, + which excludes the un-yielded message_start. A finished stream then has + exactly one usage event (message_delta) and it must be kept as-is. + """ + chunks = [ + _make_chunk(content="Yes, "), + _make_chunk(content="that works."), + _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + finish_reason="stop", + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + class TestProviderGuard: """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic @@ -297,11 +372,12 @@ class TestNonAnthropicStreamingIntact: """Make sure providers without cursor pattern still work.""" def test_completion_tokens_above_one_never_resets(self): - """Any chunk reporting completion_tokens > 1 sets saw_non_cursor - and prevents the reset.""" + """A non-Anthropic provider reporting completion_tokens > 1 from a + single usage event keeps that value.""" chunks = [ _make_chunk( - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider="openai", ), ] processor = ChunkProcessor(chunks=chunks, messages=[]) From c37d0a2e66a25a114e0af8f5d401e8b0ade1f545 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:51:04 -0700 Subject: [PATCH 69/89] fix(proxy): read general_settings without a cast and test the /team/update gate by behavior only --- .../proxy/management_endpoints/team_endpoints.py | 4 ++-- tests/test_litellm/proxy/auth/test_route_checks.py | 13 ++++--------- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f0ec3f975d5..72367b2bced 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -504,9 +504,9 @@ _GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object]) def _general_settings() -> Mapping[str, object]: - from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import general_settings - return _GENERAL_SETTINGS.validate_python(cast(object, proxy_server.general_settings)) + return _GENERAL_SETTINGS.validate_python(general_settings) def _caller_edit_access(role: TeamAccessRole | None, general_settings: Mapping[str, object]) -> TeamEditAccess: diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 2a6ba56a242..3211e85ff97 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2897,10 +2897,6 @@ def test_team_update_gate_admits_internal_user_without_org_context(): caller and update_team resolves proxy, org or team admin itself, then filters team admins through the team_admin_editable_team_fields setting. Before that the gate 401'd every team admin, which left the handler's team-admin branch unreachable.""" - from litellm.proxy._types import LiteLLMRoutes - - assert "/team/update" in LiteLLMRoutes.self_managed_routes.value - user_obj = LiteLLM_UserTable( user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -3001,11 +2997,10 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_stays_out_of_self_managed_routes(): - """Unlike POST /team/update, PATCH /team/{team_id} cannot be self-managed: its - template also matches /team/new (the collision footgun), so it stays reachable by - org admins (org_admin_allowed_routes) and proxy admins only, never by regular - internal users or through the role-agnostic self_managed_routes.""" +def test_patch_team_route_has_same_reach_as_team_update(): + """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but + NOT by regular internal users or the role-agnostic self_managed_routes — the + latter would open /team/new (the collision footgun) to any authenticated user.""" from litellm.proxy._types import LiteLLMRoutes assert RouteChecks.check_route_access( From eda81fff595f992bfae6471ef474eec896992fdc Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 15 Sep 2026 16:46:52 -0700 Subject: [PATCH 70/89] feat(ui): shared URL-state layer for tables and tabs Add useUrlTableState (search, sort, page, page size and filter_ in the query string via one nuqs useQueryStates call, with keyPrefix and urlKeys for routes that host two tables or need legacy key names) and useUrlTab (validated ?tab= param with a role-aware fallback). Migrate the Virtual Keys table onto useUrlTableState with a byte-identical URL contract and bind the Playground tab strip to ?tab=. DataTable gains controlled columnVisibility/onColumnVisibilityChange plus usePersistedColumnVisibility (localStorage per table id), and an isError prop that keeps the server page clamp from rewriting a deep-linked ?page= after a failed fetch. Virtual Keys uses both. The expired-session redirect in handleError now keeps the query string and hash so the return URL captured on re-login restores the filtered view instead of the bare list. Delete useTabRouting and tabRoutes, the pathname tab router left over from the reverted path-per-tab attempt (#34327, reverted in #34629); tab persistence has to be a query param on the static export. --- .../(dashboard)/hooks/useTabRouting.test.tsx | 82 ----- .../app/(dashboard)/hooks/useTabRouting.ts | 38 -- .../app/(dashboard)/playground/page.test.tsx | 53 ++- .../src/app/(dashboard)/playground/page.tsx | 10 +- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 152 +++++++- .../VirtualKeysPage/VirtualKeysTable.tsx | 185 ++++------ .../src/components/networking.test.ts | 38 +- .../src/components/networking.tsx | 2 +- .../shared/DataTable/DataTable.test-d.tsx | 27 +- .../shared/DataTable/DataTable.test.tsx | 98 +++++- .../components/shared/DataTable/DataTable.tsx | 15 +- .../src/components/shared/DataTable/index.ts | 2 + .../src/components/shared/DataTable/types.ts | 19 + .../usePersistedColumnVisibility.test.tsx | 98 ++++++ .../DataTable/usePersistedColumnVisibility.ts | 54 +++ .../DataTable/useUrlTableState.test.tsx | 325 ++++++++++++++++++ .../shared/DataTable/useUrlTableState.ts | 232 +++++++++++++ .../src/hooks/useUrlTab.test.tsx | 100 ++++++ ui/litellm-dashboard/src/hooks/useUrlTab.ts | 12 + .../src/utils/tabRoutes.test.ts | 47 --- ui/litellm-dashboard/src/utils/tabRoutes.ts | 26 -- 21 files changed, 1260 insertions(+), 355 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts create mode 100644 ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx create mode 100644 ui/litellm-dashboard/src/hooks/useUrlTab.ts delete mode 100644 ui/litellm-dashboard/src/utils/tabRoutes.test.ts delete mode 100644 ui/litellm-dashboard/src/utils/tabRoutes.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx deleted file mode 100644 index 24900bae798..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx +++ /dev/null @@ -1,82 +0,0 @@ -/* @vitest-environment jsdom */ -import { renderHook } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -const { mockPush, navState } = vi.hoisted(() => ({ - mockPush: vi.fn(), - navState: { pathname: "/logs" }, -})); -vi.mock("next/navigation", () => ({ - usePathname: () => navState.pathname, - useRouter: () => ({ push: mockPush }), -})); - -vi.mock("@/components/networking", () => ({ serverRootPath: "" })); - -import { createTabRoutes } from "@/utils/tabRoutes"; -import { useTabRouting } from "./useTabRouting"; - -const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const); - -const render = (ready = true) => { - const config = { - routes, - baseTabKey: "request-logs", - visibleKeys: ["audit", "deleted-keys", "deleted-teams"], - ready, - }; - return renderHook(() => useTabRouting(config)); -}; - -describe("useTabRouting", () => { - beforeEach(() => { - navState.pathname = "/logs"; - mockPush.mockClear(); - }); - - it("maps the base path to the base tab key", () => { - const { result } = render(); - expect(result.current.activeSlug).toBe(""); - expect(result.current.activeKey).toBe("request-logs"); - }); - - it("uses the slug itself as the active key for a known nested tab", () => { - navState.pathname = "/ui/logs/audit"; - const { result } = render(); - expect(result.current.activeKey).toBe("audit"); - }); - - it("falls back to the base tab key for an unknown slug", () => { - navState.pathname = "/ui/logs/bogus"; - const { result } = render(); - expect(result.current.activeKey).toBe("request-logs"); - }); - - it("redirects an unknown slug to the base href once ready", () => { - const replaceMock = vi.fn(); - const originalLocation = window.location; - Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } }); - navState.pathname = "/ui/logs/bogus"; - render(true); - expect(replaceMock).toHaveBeenCalledWith("/ui/logs/"); - Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); - }); - - it("does not redirect while not ready (role/creds still loading)", () => { - const replaceMock = vi.fn(); - const originalLocation = window.location; - Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } }); - navState.pathname = "/ui/logs/bogus"; - render(false); - expect(replaceMock).not.toHaveBeenCalled(); - Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); - }); - - it("pushes the tab href on change, mapping the base key back to the empty slug", () => { - const { result } = render(); - result.current.onTabChange("audit"); - expect(mockPush).toHaveBeenCalledWith("/ui/logs/audit/"); - result.current.onTabChange("request-logs"); - expect(mockPush).toHaveBeenCalledWith("/ui/logs/"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts deleted file mode 100644 index c17d71b4855..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts +++ /dev/null @@ -1,38 +0,0 @@ -import { useEffect } from "react"; -import { usePathname, useRouter } from "next/navigation"; -import type { TabRoutes } from "@/utils/tabRoutes"; - -interface UseTabRoutingArgs { - routes: Pick, "tabHref" | "slugFromPathname">; - baseTabKey: string; - visibleKeys: readonly string[]; - ready?: boolean; -} - -interface TabRoutingState { - activeSlug: string; - activeKey: string; - onTabChange: (key: string) => void; -} - -export function useTabRouting({ routes, baseTabKey, visibleKeys, ready = true }: UseTabRoutingArgs): TabRoutingState { - const { tabHref, slugFromPathname } = routes; - const pathname = usePathname(); - const router = useRouter(); - - const activeSlug = slugFromPathname(pathname); - const isKnownSlug = activeSlug === "" || visibleKeys.includes(activeSlug); - const activeKey = isKnownSlug ? activeSlug || baseTabKey : baseTabKey; - - useEffect(() => { - if (ready && activeSlug !== "" && !isKnownSlug) { - window.location.replace(tabHref("")); - } - }, [ready, activeSlug, isKnownSlug, tabHref]); - - const onTabChange = (key: string) => { - router.push(tabHref(key === baseTabKey ? "" : key)); - }; - - return { activeSlug, activeKey, onTabChange }; -} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx index 85e19d7d251..f12ebc0b831 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx @@ -1,5 +1,8 @@ -import { render, screen } from "@testing-library/react"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; import PlaygroundPage from "./page"; const authState = { userRole: "Admin" }; @@ -35,14 +38,17 @@ vi.mock("@/app/(dashboard)/playground/components/chat_ui/AgentBuilderView", () = default: () =>
, })); -describe("PlaygroundPage role guard", () => { - beforeEach(() => { - authState.userRole = "Admin"; - }); +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + onUrlUpdate.mock.calls.at(-1)?.[0]; +beforeEach(() => { + authState.userRole = "Admin"; +}); + +describe("PlaygroundPage role guard", () => { it.each(["Internal Viewer", "Admin Viewer"])("blocks the entire playground for %s", (role) => { authState.userRole = role; - render(); + renderWithProviders(); expect(screen.getByText("Access Denied")).toBeInTheDocument(); expect(screen.queryByRole("tab")).not.toBeInTheDocument(); @@ -54,10 +60,43 @@ describe("PlaygroundPage role guard", () => { it.each(["Admin", "Internal User", "Org Admin"])("renders the playground for %s", (role) => { authState.userRole = role; - render(); + renderWithProviders(); expect(screen.queryByText("Access Denied")).not.toBeInTheDocument(); expect(screen.getByRole("tab", { name: "Chat" })).toBeInTheDocument(); expect(screen.getByTestId("chat-ui")).toBeInTheDocument(); }); }); + +describe("PlaygroundPage ?tab= deep link", () => { + it("opens on Chat when the URL has no tab", () => { + renderWithProviders(); + + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true"); + }); + + it("activates the tab named in ?tab=", () => { + renderWithProviders(, { searchParams: { tab: "compare" } }); + + expect(screen.getByRole("tab", { name: "Compare" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "false"); + }); + + it("falls back to Chat when ?tab= is not a playground tab", () => { + renderWithProviders(, { searchParams: { tab: "settings" } }); + + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true"); + }); + + it("clicking a tab writes ?tab= with history replace", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + await user.click(screen.getByRole("tab", { name: "Compliance" })); + + expect(await screen.findByRole("tab", { name: "Compliance", selected: true })).toBeInTheDocument(); + await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compliance")); + expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index 78ca538d8b5..27a61415672 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -9,6 +9,9 @@ import { DeprecationBanner } from "@/components/DeprecationBanner"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchProxySettings } from "@/utils/proxyUtils"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { useUrlTab } from "@/hooks/useUrlTab"; + +const PLAYGROUND_TABS = ["chat", "compare", "compliance", "agent-builder"] as const; interface ProxySettings { PROXY_BASE_URL?: string; @@ -18,6 +21,7 @@ interface ProxySettings { export default function PlaygroundPage() { const { accessToken, userRole, userId, disabledPersonalKeyCreation, token, isViewOnly } = useAuthorized(); const [proxySettings, setProxySettings] = useState(undefined); + const [activeTab, setActiveTab] = useUrlTab(PLAYGROUND_TABS, "chat"); useEffect(() => { const initializeProxySettings = async () => { @@ -48,7 +52,11 @@ export default function PlaygroundPage() { return (
- + Chat diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 6c742309eb7..4d963a2f603 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -4,7 +4,7 @@ import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest"; import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; -import { KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; +import { KEY_TABLE_HIDDEN_COLUMNS, KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; @@ -187,6 +187,7 @@ const lastHistoryMode = (onUrlUpdate: Mock) => onUrlUpdate. beforeEach(() => { vi.clearAllMocks(); + localStorage.clear(); mockUseKeys.mockReturnValue(keysResult([mockKey])); mockUseKeyInfo.mockReturnValue(keyInfoResult(undefined)); @@ -823,16 +824,27 @@ describe("table state lives in the URL so it survives leaving and returning to t }); it("restores the drawer filters from the URL on mount", async () => { - renderWithProviders(, { searchParams: { filter_team: "team-1", filter_user: "user-42" } }); + const searchParams = { + filter_team: "team-1", + filter_org: "org-1", + filter_user: "user-42", + filter_key_id: mockKey.token, + }; + const expectedKeyListOptions = { + teamID: "team-1", + organizationID: "org-1", + userID: "user-42", + keyHash: mockKey.token, + }; + renderWithProviders(, { searchParams }); await waitFor(() => { - expect(mockUseKeys).toHaveBeenLastCalledWith( - 1, - 50, - expect.objectContaining({ teamID: "team-1", userID: "user-42" }), - ); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining(expectedKeyListOptions)); }); expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team"); + expect(screen.getByTestId("filter-chip-org_id")).toHaveTextContent("Test Organization"); + expect(screen.getByTestId("filter-chip-user_id")).toHaveTextContent("user-42"); + expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent(mockKey.token); }); it("restores the status filter from the URL and sends it to /key/list", async () => { @@ -853,6 +865,21 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.queryByTestId("filter-chip-status")).not.toBeInTheDocument(); }); + it("drops a hand-edited status from the URL when another filter chip is removed", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: { filter_status: "bogus", filter_user: "user-42" }, + onUrlUpdate, + }); + + fireEvent.click(await screen.findByTestId("filter-chip-remove-user_id")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_user")).toBeNull(); + }); + expect(lastSearchParam(onUrlUpdate, "filter_status")).toBeNull(); + }); + it("writes the search term to the URL", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { onUrlUpdate }); @@ -896,6 +923,40 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.queryByTestId("filter-chip-user_id")).not.toBeInTheDocument(); }); + it("writes the Organization and Key ID drawer filters to the URL and clears them again", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + openFilters(); + await chooseSelectOption(user, await screen.findByPlaceholderText(/Select an organization/), /Test Organization/); + fireEvent.change(screen.getByPlaceholderText(/Enter Key ID/), { target: { value: mockKey.token } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_org")).toBe("org-1"); + }); + expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBe(mockKey.token); + expect(lastSearchParam(onUrlUpdate, "filter_org_id")).toBeNull(); + expect(lastSearchParam(onUrlUpdate, "filter_key_hash")).toBeNull(); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ organizationID: "org-1", keyHash: mockKey.token }), + ); + }); + + fireEvent.click(screen.getByTestId("datatable-clear-filters")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_org")).toBeNull(); + }); + expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBeNull(); + expect(screen.queryByTestId("filter-chip-org_id")).not.toBeInTheDocument(); + expect(screen.queryByTestId("filter-chip-key_hash")).not.toBeInTheDocument(); + }); + it("returns to page 1 when the search term changes", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { searchParams: { page: "3" }, onUrlUpdate }); @@ -953,14 +1014,16 @@ describe("table state lives in the URL so it survives leaving and returning to t }); }); - it("falls back to the default sort when the URL names a column the table cannot sort by", async () => { - renderWithProviders(, { searchParams: { sort_by: "totally_unknown_field" } }); + it("falls back to the default sort column, keeping the URL's direction, when the table cannot sort by sort_by", async () => { + renderWithProviders(, { + searchParams: { sort_by: "totally_unknown_field", sort_order: "asc" }, + }); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith( 1, 50, - expect.objectContaining({ sortBy: "created_at", sortOrder: "desc" }), + expect.objectContaining({ sortBy: "created_at", sortOrder: "asc" }), ); }); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); @@ -1003,3 +1066,72 @@ describe("table state lives in the URL so it survives leaving and returning to t }); }); }); + +describe("column choices survive a reload", () => { + const STORAGE_KEY = "litellm_table_columns_virtual-keys"; + const storedColumns = () => JSON.parse(localStorage.getItem(STORAGE_KEY) ?? "null"); + + it("hides a column that was hidden on a previous visit while the default-hidden columns stay hidden", () => { + localStorage.setItem(STORAGE_KEY, JSON.stringify({ budget_reset_at: false })); + + renderWithProviders(); + + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.queryByText("Budget Reset")).not.toBeInTheDocument(); + expect(screen.queryByText("Created By")).not.toBeInTheDocument(); + }); + + it("writes a column toggled on through the Columns menu to storage and shows it again on the next mount", async () => { + const user = userEvent.setup(); + const { unmount } = renderWithProviders(); + expect(screen.queryByText("Created By")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Created By")); + await user.keyboard("{Escape}"); + + expect(storedColumns()).toEqual({ ...KEY_TABLE_HIDDEN_COLUMNS, created_by: true }); + + unmount(); + renderWithProviders(); + + expect(screen.getByText("Created By")).toBeInTheDocument(); + }); +}); + +describe("a failed keys fetch does not rewrite the URL", () => { + const renderOnPage3OfMany = async () => { + mockUseKeys.mockReturnValue(keysResult([mockKey], { total_count: 200, total_pages: 4 })); + const onUrlUpdate = vi.fn(); + const view = renderWithProviders(, { searchParams: { page: "3" }, onUrlUpdate }); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything()); + }); + return { ...view, onUrlUpdate }; + }; + + it("keeps ?page=3 when the keys query errors, instead of snapping to page 1 on the empty count", async () => { + const { rerender, onUrlUpdate } = await renderOnPage3OfMany(); + + mockUseKeys.mockReturnValue(keysResult([], {}, { data: undefined, isError: true })); + rerender(); + + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything()); + expect(onUrlUpdate).not.toHaveBeenCalled(); + }); + + it("still snaps ?page=3 back to the first page when the keys query succeeds with no rows", async () => { + const { rerender, onUrlUpdate } = await renderOnPage3OfMany(); + + mockUseKeys.mockReturnValue(keysResult([])); + rerender(); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.anything()); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 1f52bdd7335..39dd2cc5ab2 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -10,15 +10,18 @@ import { DataTableFilterDrawer, DataTableFilterField, DataTableToolbar, + usePersistedColumnVisibility, + useUrlTableState, + type UrlTableStateOptions, } from "@/components/shared/DataTable"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { PageHeader } from "@/components/shared/PageHeader"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; -import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { ColumnFiltersState, functionalUpdate, OnChangeFn } from "@tanstack/react-table"; import { KeyRound } from "lucide-react"; -import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryState, useQueryStates } from "nuqs"; +import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; @@ -56,44 +59,30 @@ const STATUS_FILTER_ITEMS = [ ...KEY_STATUS_VALUES.map((value) => ({ value, label: KEY_STATUS_LABELS[value] })), ]; -const isKeyStatusFilter = (value: string): value is KeyStatusFilter => - (KEY_STATUS_VALUES as readonly string[]).includes(value); +const isKeyStatusFilter = (value: unknown): value is KeyStatusFilter => + (KEY_STATUS_VALUES as readonly unknown[]).includes(value); -const DEFAULT_SORT_BY = "created_at"; -const DEFAULT_SORT_ORDER = "desc"; -const DEFAULT_PAGE_SIZE = 50; -const MAX_PAGE_SIZE = 100; -const MAX_PAGE = 100_000; +const isUsableFilter = (filter: ColumnFiltersState[number]): boolean => + filter.id !== "status" || isKeyStatusFilter(filter.value); -const boundedInteger = (min: number, max: number, fallback: number) => - createParser({ - parse: (value: string) => { - const parsed = parseAsInteger.parse(value); - return parsed === null ? null : Math.min(Math.max(parsed, min), max); - }, - serialize: String, - }).withDefault(fallback); - -// The filters carry a prefix because /api-keys also takes team_id, key_alias and key_type -// as create-key prefills; an unprefixed filter would hijack those deep links. -const TABLE_STATE = { - key_search: parseAsString.withDefault(""), - sort_by: parseAsString.withDefault(DEFAULT_SORT_BY), - sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault(DEFAULT_SORT_ORDER), - page: boundedInteger(1, MAX_PAGE, 1), - page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE), - filter_team: parseAsString.withDefault(""), - filter_org: parseAsString.withDefault(""), - filter_user: parseAsString.withDefault(""), - filter_key_id: parseAsString.withDefault(""), - filter_status: parseAsString.withDefault(""), +const TABLE_STATE_OPTIONS: UrlTableStateOptions = { + sortFields: KEY_TABLE_SORT_FIELDS, + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: 50, + maxPageSize: 100, + filterColumns: FILTER_COLUMNS, + urlKeys: { + search: "key_search", + filter_team_id: "filter_team", + filter_org_id: "filter_org", + filter_user_id: "filter_user", + filter_key_hash: "filter_key_id", + }, }; -const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc"); - -const filterValue = (filters: ColumnFiltersState, column: FilterColumn): string | null => { +const appliedFilter = (filters: ColumnFiltersState, column: FilterColumn): string | undefined => { const value = filters.find((filter) => filter.id === column)?.value; - return (typeof value === "string" ? value.trim() : "") || null; + return typeof value === "string" ? value : undefined; }; export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { @@ -103,50 +92,38 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); const [selectedKeyId, setSelectedKeyId] = useQueryState("key", parseAsString.withOptions({ history: "push" })); - const [tableState, setTableState] = useQueryStates(TABLE_STATE); + const { + search: searchInput, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters: urlColumnFilters, + onColumnFiltersChange: setUrlColumnFilters, + } = useUrlTableState(TABLE_STATE_OPTIONS); + const columnFilters = useMemo(() => urlColumnFilters.filter(isUsableFilter), [urlColumnFilters]); + const onColumnFiltersChange = useCallback>( + (updaterOrValue) => setUrlColumnFilters(functionalUpdate(updaterOrValue, columnFilters)), + [columnFilters, setUrlColumnFilters], + ); + const { columnVisibility, onColumnVisibilityChange } = usePersistedColumnVisibility( + "virtual-keys", + KEY_TABLE_HIDDEN_COLUMNS, + ); const [filtersOpen, setFiltersOpen] = useState(false); - const searchInput = tableState.key_search; const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); - // A hand-edited sort_by the table cannot sort by would 400 at /key/list and leave the page loading. - const sortBy = KEY_TABLE_SORT_FIELDS.includes(tableState.sort_by) ? tableState.sort_by : DEFAULT_SORT_BY; - const sorting = useMemo( - () => [{ id: sortBy, desc: tableState.sort_order === "desc" }], - [sortBy, tableState.sort_order], - ); - const tablePagination = useMemo( - () => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }), - [tableState.page, tableState.page_size], - ); - const { filter_team, filter_org, filter_user, filter_key_id, filter_status } = tableState; - const appliedFilters = useMemo( - () => ({ - team_id: filter_team.trim(), - org_id: filter_org.trim(), - user_id: filter_user.trim(), - key_hash: filter_key_id.trim(), - status: isKeyStatusFilter(filter_status) ? filter_status : "", - }), - [filter_team, filter_org, filter_user, filter_key_id, filter_status], - ); - const columnFilters = useMemo( - () => - FILTER_COLUMNS.filter((column) => appliedFilters[column]).map((column) => ({ - id: column, - value: appliedFilters[column], - })), - [appliedFilters], - ); - + const [activeSort] = sorting; const keyListOptions = { - teamID: appliedFilters.team_id || undefined, - organizationID: appliedFilters.org_id || undefined, + teamID: appliedFilter(columnFilters, "team_id"), + organizationID: appliedFilter(columnFilters, "org_id"), search: searchQuery.trim() || undefined, - userID: appliedFilters.user_id || undefined, - keyHash: appliedFilters.key_hash || undefined, - status: appliedFilters.status || undefined, - sortBy, - sortOrder: tableState.sort_order, + userID: appliedFilter(columnFilters, "user_id"), + keyHash: appliedFilter(columnFilters, "key_hash"), + status: appliedFilter(columnFilters, "status"), + sortBy: activeSort.id, + sortOrder: activeSort.desc ? "desc" : "asc", expand: "user", }; @@ -155,55 +132,13 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { isPending, isPlaceholderData, isFetching, + isError, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); + } = useKeys(pagination.pageIndex + 1, pagination.pageSize, keyListOptions); const keyList = useMemo(() => keys?.keys ?? [], [keys]); const rowCount = keys?.total_count ?? 0; - const handleSearchChange = useCallback( - (value: string) => { - void setTableState({ key_search: value || null, page: null }); - }, - [setTableState], - ); - - const handleSortingChange = useCallback>( - (updaterOrValue) => { - const active = functionalUpdate(updaterOrValue, sorting)[0]; - void setTableState({ - sort_by: active?.id ?? null, - sort_order: active ? toSortOrder(active) : null, - page: null, - }); - }, - [sorting, setTableState], - ); - - const handleColumnFiltersChange = useCallback>( - (updaterOrValue) => { - const next = functionalUpdate(updaterOrValue, columnFilters); - const nextFilters = { - filter_team: filterValue(next, "team_id"), - filter_org: filterValue(next, "org_id"), - filter_user: filterValue(next, "user_id"), - filter_key_id: filterValue(next, "key_hash"), - filter_status: filterValue(next, "status"), - page: null, - }; - void setTableState(nextFilters); - }, - [columnFilters, setTableState], - ); - - const handlePaginationChange = useCallback>( - (updaterOrValue) => { - const next = functionalUpdate(updaterOrValue, tablePagination); - void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize }); - }, - [tablePagination, setTableState], - ); - const columns = useMemo( () => getKeyTableColumns({ allTeams, organizations, onSelectKey: (key) => void setSelectedKeyId(key.token) }), [allTeams, organizations, setSelectedKeyId], @@ -296,20 +231,22 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { data={keyList} columns={columns} getRowId={(row) => row.token} - defaultColumnVisibility={KEY_TABLE_HIDDEN_COLUMNS} + columnVisibility={columnVisibility} + onColumnVisibilityChange={onColumnVisibilityChange} sortingMode="server" sorting={sorting} - onSortingChange={handleSortingChange} + onSortingChange={onSortingChange} paginationMode="server" - pagination={tablePagination} - onPaginationChange={handlePaginationChange} + pagination={pagination} + onPaginationChange={onPaginationChange} rowCount={rowCount} filterMode="server" columnFilters={columnFilters} - onColumnFiltersChange={handleColumnFiltersChange} + onColumnFiltersChange={onColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" isLoading={isPending || isPlaceholderData} + isError={isError} loadingMessage="Loading keys..." noDataMessage="No keys found" fillHeight @@ -319,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { refetch?.()} isRefreshing={isFetching} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 578e355b85d..3b2a17101ee 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -20,25 +20,39 @@ describe("networking - expired session handling", () => { global.fetch = originalFetch; }); - it("should call clearTokenCookies on expired session", async () => { - const errorData = "Authentication Error - Expired Key"; - const { toast } = await import("@/lib/toast"); + const loadFreshHandleError = async () => { + vi.resetModules(); + const fresh = await import("./networking"); + return fresh.handleError; + }; - if (errorData.includes("Authentication Error - Expired Key")) { - toast.info("UI Session Expired. Logging out."); - clearTokenCookies(); - } + const stubLocation = (pathname: string, search: string, hash: string) => { + const location = { pathname, search, hash, href: "" }; + vi.stubGlobal("window", { location }); + return location; + }; + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("keeps the query string and hash on the redirect after session expiry", async () => { + const handleError = await loadFreshHandleError(); + const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", "#row-3"); + + await handleError("Authentication Error - Expired Key"); + + expect(location.href).toBe("/ui/api-keys/?filter_team=t1&page=2#row-3"); expect(clearTokenCookies).toHaveBeenCalledOnce(); }); - it("should not clear cookies for non-authentication errors", () => { - const errorData = "Some other error"; + it("does not navigate or clear cookies for other errors", async () => { + const handleError = await loadFreshHandleError(); + const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", ""); - if (errorData.includes("Authentication Error - Expired Key")) { - clearTokenCookies(); - } + await handleError("Some other error"); + expect(location.href).toBe(""); expect(clearTokenCookies).not.toHaveBeenCalled(); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index cab073dc808..e77c8ba7e41 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -383,7 +383,7 @@ export const handleError = async (errorData: string | any) => { clearTokenCookies(); const browserLocation = getWindowLocation(); if (browserLocation) { - window.location.href = browserLocation.pathname; + window.location.href = browserLocation.pathname + browserLocation.search + browserLocation.hash; } } lastErrorTime = currentTime; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx index 7bbd4f918cd..c6c6a392aec 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx @@ -1,4 +1,10 @@ -import type { ColumnDef, PaginationState, RowSelectionState, SortingState } from "@tanstack/react-table"; +import type { + ColumnDef, + PaginationState, + RowSelectionState, + SortingState, + VisibilityState, +} from "@tanstack/react-table"; import { DataTable } from "./DataTable"; @@ -12,6 +18,7 @@ const columns: ColumnDef[] = []; const sorting: SortingState = [{ id: "name", desc: false }]; const pagination: PaginationState = { pageIndex: 0, pageSize: 10 }; const rowSelection: RowSelectionState = { r1: true }; +const columnVisibility: VisibilityState = { name: false }; const noop = () => {}; export const uncontrolled = ; @@ -32,6 +39,8 @@ export const controlled = ( onColumnFiltersChange={noop} rowSelection={rowSelection} onRowSelectionChange={noop} + columnVisibility={columnVisibility} + onColumnVisibilityChange={noop} /> ); @@ -65,3 +74,19 @@ export const selectionWithoutHandler = ( // @ts-expect-error a controlled `rowSelection` needs `onRowSelectionChange` or selection changes are dropped ); + +export const visibilityWithoutHandler = ( + // @ts-expect-error a controlled `columnVisibility` needs `onColumnVisibilityChange` or Columns-menu toggles are dropped + +); + +export const bothVisibilitySources = ( + // @ts-expect-error `defaultColumnVisibility` seeds uncontrolled visibility, so it cannot pair with a controlled `columnVisibility` + +); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 8ed8e392ae1..336fc43d695 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -1,4 +1,4 @@ -import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState } from "@tanstack/react-table"; +import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState, VisibilityState } from "@tanstack/react-table"; import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; @@ -278,11 +278,18 @@ describe("DataTable pagination", () => { type ServerPageHarnessProps = { rowCount: number; isLoading?: boolean; + isError?: boolean; initialPageIndex: number; onChange: (next: PaginationState) => void; }; - function ServerPageHarness({ rowCount, isLoading = false, initialPageIndex, onChange }: ServerPageHarnessProps) { + function ServerPageHarness({ + rowCount, + isLoading = false, + isError = false, + initialPageIndex, + onChange, + }: ServerPageHarnessProps) { const [pagination, setPagination] = useState({ pageIndex: initialPageIndex, pageSize: 10 }); const handleChange: OnChangeFn = (updater) => { const next = typeof updater === "function" ? updater(pagination) : updater; @@ -298,6 +305,7 @@ describe("DataTable pagination", () => { onPaginationChange={handleChange} rowCount={rowCount} isLoading={isLoading} + isError={isError} /> ); } @@ -339,6 +347,29 @@ describe("DataTable pagination", () => { expect(onChange).toHaveBeenCalledTimes(1); expect(screen.getByText("Page 2 of 2")).toBeInTheDocument(); }); + + it("server mode keeps a deep-linked page when the fetch failed, instead of snapping to page 1 on rowCount 0", async () => { + const onChange = vi.fn(); + render(); + + expect(screen.getByText("Page 3 of 1")).toBeInTheDocument(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => { + const onChange = vi.fn(); + const { rerender } = render(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + + rerender(); + + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 1, pageSize: 10 })); + expect(onChange).toHaveBeenCalledTimes(1); + expect(screen.getByText("Page 2 of 2")).toBeInTheDocument(); + }); }); describe("DataTable filtering", () => { @@ -555,6 +586,69 @@ describe("DataTable column visibility", () => { expect(await screen.findByTestId("view-option-email")).toBeInTheDocument(); expect(screen.queryByTestId("view-option-name")).not.toBeInTheDocument(); }); + + it("uncontrolled mode seeds hidden columns from defaultColumnVisibility and still toggles internally", async () => { + const user = userEvent.setup(); + render( + } + />, + ); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); + }); + + it("controlled mode hides columns from the prop and reports toggles without changing them locally", async () => { + const user = userEvent.setup(); + const onColumnVisibilityChange = vi.fn>(); + render( + } + />, + ); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + + expect(onColumnVisibilityChange).toHaveBeenCalledTimes(1); + const updater = onColumnVisibilityChange.mock.calls[0]?.[0]; + const next = typeof updater === "function" ? updater({ email: false }) : updater; + expect(next).toEqual({ email: true }); + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + }); + + it("controlled mode reveals the column once the parent applies the reported change", async () => { + const user = userEvent.setup(); + const Harness = () => { + const [columnVisibility, setColumnVisibility] = useState({ email: false }); + return ( + } + /> + ); + }; + render(); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); + }); }); describe("DataTable pinned columns", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 26162a3f1f7..e0f57ae1052 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -457,6 +457,7 @@ function useDataTableInstance( onPaginationChange, rowCount, isLoading = false, + isError, pageSizeOptions = DEFAULT_PAGE_SIZE_OPTIONS, filterMode = "none", columnFilters, @@ -466,6 +467,8 @@ function useDataTableInstance( onGlobalFilterChange, enableColumnResizing = false, columnResizeMode = "onEnd", + columnVisibility, + onColumnVisibilityChange, defaultColumnVisibility, getRowCanExpand, renderSubComponent, @@ -481,7 +484,7 @@ function useDataTableInstance( pageIndex: 0, pageSize: pageSizeOptions[0] ?? 25, }); - useServerPageClamp(paginationMode === "server" && !isLoading, rowCount, paginationState); + useServerPageClamp(paginationMode === "server" && !isLoading && !isError, rowCount, paginationState); const filterState = useControllable( columnFilters, onColumnFiltersChange, @@ -490,7 +493,11 @@ function useDataTableInstance( const globalFilterState = useControllable(globalFilter, onGlobalFilterChange, ""); const expandedState = useControllable(expanded, onExpandedChange, {}); const rowSelectionState = useControllable(rowSelection, onRowSelectionChange, {}); - const [columnVisibility, setColumnVisibility] = useState(defaultColumnVisibility ?? {}); + const columnVisibilityState = useControllable( + columnVisibility, + onColumnVisibilityChange, + defaultColumnVisibility ?? {}, + ); const [columnSizing, setColumnSizing] = useState({}); const columnPinning = React.useMemo(() => derivePinning(columns), [columns]); const expansionGuard = renderSubComponent !== undefined ? getRowCanExpand : undefined; @@ -505,7 +512,7 @@ function useDataTableInstance( globalFilter: globalFilterState.value, expanded: expandedState.value, rowSelection: rowSelectionState.value, - columnVisibility, + columnVisibility: columnVisibilityState.value, columnSizing, }, initialState: { columnPinning }, @@ -521,7 +528,7 @@ function useDataTableInstance( onGlobalFilterChange: globalFilterState.onChange, onExpandedChange: expandedState.onChange, onRowSelectionChange: rowSelectionState.onChange, - onColumnVisibilityChange: setColumnVisibility, + onColumnVisibilityChange: columnVisibilityState.onChange, onColumnSizingChange: setColumnSizing, getColumnCanGlobalFilter: (column) => columnCanGlobalFilter(data[0], column), getCoreRowModel: getCoreRowModel(), diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index 39a887ba948..85cc5f287e0 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -12,6 +12,8 @@ export { type DataTableSortVariant, type DataTableSortField, } from "./DataTableSortHeader"; +export { usePersistedColumnVisibility } from "./usePersistedColumnVisibility"; +export { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "./useUrlTableState"; export type { DataTablePaginationProps } from "./DataTablePagination"; export type { ColumnPinnedSide, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index c767a0a64c0..4529e3df164 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -27,6 +27,7 @@ export interface DataTableResolvedProps { getRowId?: (row: TData, index: number, parent?: Row) => string; isLoading?: boolean; + isError?: boolean; loadingMessage?: string; skeletonRowCount?: number; noDataMessage?: React.ReactNode; @@ -53,6 +54,8 @@ export interface DataTableResolvedProps { enableColumnResizing?: boolean; columnResizeMode?: ColumnResizeMode; + columnVisibility?: VisibilityState; + onColumnVisibilityChange?: OnChangeFn; defaultColumnVisibility?: VisibilityState; getRowCanExpand?: (row: Row) => boolean; @@ -96,6 +99,9 @@ type DataTableBaseProps = Omit< | "columnFilters" | "onColumnFiltersChange" | "defaultColumnFilters" + | "columnVisibility" + | "onColumnVisibilityChange" + | "defaultColumnVisibility" | "rowSelection" | "onRowSelectionChange" >; @@ -142,6 +148,18 @@ type FilterProps = defaultColumnFilters?: ColumnFiltersState; }; +type ColumnVisibilityProps = + | { + columnVisibility: VisibilityState; + onColumnVisibilityChange: OnChangeFn; + defaultColumnVisibility?: never; + } + | { + columnVisibility?: never; + onColumnVisibilityChange?: never; + defaultColumnVisibility?: VisibilityState; + }; + type RowSelectionProps = | { rowSelection: RowSelectionState; onRowSelectionChange: OnChangeFn } | { rowSelection?: never; onRowSelectionChange?: OnChangeFn }; @@ -150,4 +168,5 @@ export type DataTableProps = DataTableBaseProps `litellm_table_columns_${tableId}`; + +const stored = (tableId: string): unknown => { + const raw = localStorage.getItem(keyFor(tableId)); + return raw === null ? null : JSON.parse(raw); +}; + +describe("usePersistedColumnVisibility", () => { + beforeEach(() => { + localStorage.clear(); + }); + + afterEach(() => { + localStorage.clear(); + vi.restoreAllMocks(); + }); + + it("layers the stored choices over the defaults, so a default added after the snapshot still applies", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false, spend: true })); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false, name: false })); + + expect(result.current.columnVisibility).toEqual({ email: false, spend: true, name: false }); + }); + + it("falls back to the defaults when nothing is stored, and to {} without defaults", () => { + const withDefaults = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + expect(withDefaults.result.current.columnVisibility).toEqual({ spend: false }); + + const bare = renderHook(() => usePersistedColumnVisibility("keys")); + expect(bare.result.current.columnVisibility).toEqual({}); + }); + + it("writes an object update to state and storage", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + + expect(result.current.columnVisibility).toEqual({ email: false }); + expect(stored("keys")).toEqual({ email: false }); + }); + + it("resolves a function updater against the current state before persisting", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false }))); + + expect(result.current.columnVisibility).toEqual({ email: false, name: false }); + expect(stored("keys")).toEqual({ email: false, name: false }); + }); + + it.each([ + ["truncated JSON", '{"email":fal'], + ["a JSON scalar", "42"], + ["a JSON array", "[true]"], + ["non-boolean values", JSON.stringify({ email: "no" })], + ])("falls back to the defaults when storage holds %s", (_label, raw) => { + localStorage.setItem(keyFor("keys"), raw); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + + expect(result.current.columnVisibility).toEqual({ spend: false }); + }); + + it("keeps distinct tableIds isolated in state and storage", () => { + const keys = renderHook(() => usePersistedColumnVisibility("keys")); + const teams = renderHook(() => usePersistedColumnVisibility("teams")); + + act(() => keys.result.current.onColumnVisibilityChange({ email: false })); + + expect(keys.result.current.columnVisibility).toEqual({ email: false }); + expect(teams.result.current.columnVisibility).toEqual({}); + expect(stored("keys")).toEqual({ email: false }); + expect(stored("teams")).toBeNull(); + }); + + it("returns the defaults without throwing when storage is unavailable", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { + throw new Error("SecurityError"); + }); + vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { + throw new Error("QuotaExceededError"); + }); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + expect(result.current.columnVisibility).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + expect(result.current.columnVisibility).toEqual({ email: false }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts new file mode 100644 index 00000000000..b56ae13d63b --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -0,0 +1,54 @@ +import type { OnChangeFn, VisibilityState } from "@tanstack/react-table"; +import { useCallback, useState } from "react"; + +import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; + +const STORAGE_KEY_PREFIX = "litellm_table_columns_"; + +const EMPTY_VISIBILITY: VisibilityState = {}; + +function storageKey(tableId: string): string { + return `${STORAGE_KEY_PREFIX}${tableId}`; +} + +function isVisibilityState(value: unknown): value is VisibilityState { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + return false; + } + return Object.values(value).every((visible) => typeof visible === "boolean"); +} + +function readStoredVisibility(tableId: string, defaults: VisibilityState): VisibilityState { + const raw = getLocalStorageItem(storageKey(tableId)); + if (raw === null) { + return defaults; + } + try { + const parsed: unknown = JSON.parse(raw); + return isVisibilityState(parsed) ? { ...defaults, ...parsed } : defaults; + } catch { + return defaults; + } +} + +export function usePersistedColumnVisibility( + tableId: string, + defaults: VisibilityState = EMPTY_VISIBILITY, +): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn } { + const [columnVisibility, setColumnVisibility] = useState(() => + readStoredVisibility(tableId, defaults), + ); + + const onColumnVisibilityChange = useCallback>( + (updater) => { + setColumnVisibility((previous) => { + const next = typeof updater === "function" ? updater(previous) : updater; + setLocalStorageItem(storageKey(tableId), JSON.stringify(next)); + return next; + }); + }, + [tableId], + ); + + return { columnVisibility, onColumnVisibilityChange }; +} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx new file mode 100644 index 00000000000..ad46d18b5d8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx @@ -0,0 +1,325 @@ +import { SortingState } from "@tanstack/react-table"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { withNuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import { describe, expect, it, Mock, vi } from "vitest"; +import { useUrlTableState, type UrlTableStateOptions } from "./useUrlTableState"; + +const FILTER_COLUMNS = ["team_id", "user_id"] as const; +type FilterColumn = (typeof FILTER_COLUMNS)[number]; + +const BASE_OPTIONS: UrlTableStateOptions = { + sortFields: ["created_at", "spend", "key_alias"], + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: 50, + filterColumns: FILTER_COLUMNS, +}; + +const PREFIXED_AND_UNPREFIXED_PARAMS = { + audit_page: "2", + audit_page_size: "10", + audit_search: "prefixed", + audit_sort_by: "spend", + audit_sort_order: "asc", + audit_filter_team_id: "team-1", + page: "5", + search: "unprefixed", + filter_team_id: "other-team", +}; + +const RENAMED_AND_DEFAULT_PARAMS = { + key_search: "prod", + filter_team: "team-1", + search: "ignored", + filter_team_id: "ignored", +}; + +const flipDirection = (previous: SortingState): SortingState => previous.map((sort) => ({ ...sort, desc: !sort.desc })); + +const renderTableState = ( + searchParams: Record = {}, + overrides: Partial> = {}, +) => { + const onUrlUpdate = vi.fn(); + const options = { ...BASE_OPTIONS, ...overrides }; + const hook = renderHook(() => useUrlTableState(options), { + wrapper: withNuqsTestingAdapter({ searchParams, onUrlUpdate, hasMemory: true }), + }); + return { ...hook, onUrlUpdate }; +}; + +const lastUrl = (onUrlUpdate: Mock) => { + const event = onUrlUpdate.mock.calls.at(-1)?.[0]; + if (!event) throw new Error("no URL update was emitted"); + return event; +}; + +const flushUrl = async (onUrlUpdate: Mock, write: () => void) => { + const callsBefore = onUrlUpdate.mock.calls.length; + await act(async () => { + write(); + }); + await waitFor(() => expect(onUrlUpdate.mock.calls.length).toBeGreaterThan(callsBefore)); + return lastUrl(onUrlUpdate).searchParams; +}; + +describe("reading table state from the URL", () => { + it("falls back to the defaults when the URL carries no table state", () => { + const { result } = renderTableState(); + + expect(result.current.search).toBe(""); + expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]); + expect(result.current.pagination).toEqual({ pageIndex: 0, pageSize: 50 }); + expect(result.current.columnFilters).toEqual([]); + }); + + it("maps the 1-based page and page_size onto TanStack pagination", () => { + const { result } = renderTableState({ page: "3", page_size: "25" }); + + expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 25 }); + }); + + it.each(["0", "-3", "not-a-number"])("clamps a page of %s up to the first page", (page) => { + const { result } = renderTableState({ page }); + + expect(result.current.pagination.pageIndex).toBe(0); + }); + + it.each([ + ["1000", undefined, 100], + ["1000", 20, 20], + ["0", undefined, 1], + ])("clamps a page_size of %s with maxPageSize %s to %s", (pageSize, maxPageSize, expected) => { + const { result } = renderTableState({ page_size: pageSize }, { maxPageSize }); + + expect(result.current.pagination.pageSize).toBe(expected); + }); + + it("reads a sortable sort_by and its sort_order", () => { + const { result } = renderTableState({ sort_by: "spend", sort_order: "asc" }); + + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("resolves a sort_by outside the allow-list to the default column while keeping the URL's direction", () => { + const { result } = renderTableState({ sort_by: "totally_unknown", sort_order: "asc" }); + + expect(result.current.sorting).toEqual([{ id: "created_at", desc: false }]); + }); + + it("maps filter_ params onto columnFilters, trimming whitespace and dropping blanks", () => { + const { result } = renderTableState({ filter_team_id: "team-1", filter_user_id: " " }); + + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + + const trimmed = renderTableState({ filter_user_id: " user-42 " }); + expect(trimmed.result.current.columnFilters).toEqual([{ id: "user_id", value: "user-42" }]); + }); + + it("reads the search term verbatim so the input can hold trailing spaces", () => { + const { result } = renderTableState({ search: "prod " }); + + expect(result.current.search).toBe("prod "); + }); + + it("reads every key under keyPrefix and ignores the unprefixed ones", () => { + const { result } = renderTableState(PREFIXED_AND_UNPREFIXED_PARAMS, { keyPrefix: "audit_" }); + + expect(result.current.pagination).toEqual({ pageIndex: 1, pageSize: 10 }); + expect(result.current.search).toBe("prefixed"); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("reads renamed keys from urlKeys and ignores the default names", () => { + const { result } = renderTableState(RENAMED_AND_DEFAULT_PARAMS, { + urlKeys: { search: "key_search", filter_team_id: "filter_team" }, + }); + + expect(result.current.search).toBe("prod"); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("applies keyPrefix in front of a renamed key", () => { + const { result } = renderTableState( + { audit_key_search: "prod", key_search: "ignored" }, + { keyPrefix: "audit_", urlKeys: { search: "key_search" } }, + ); + + expect(result.current.search).toBe("prod"); + }); +}); + +describe("writing table state to the URL", () => { + it("resolves a function updater against the current pagination and replaces history", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "2" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange((previous) => ({ ...previous, pageIndex: previous.pageIndex + 1 })), + ); + + expect(url.get("page")).toBe("3"); + expect(url.has("page_size")).toBe(false); + expect(lastUrl(onUrlUpdate).options.history).toBe("replace"); + expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 50 }); + }); + + it("writes page_size and drops it again once it returns to the default", async () => { + const { result, onUrlUpdate } = renderTableState(); + + const withSize = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange({ pageIndex: 0, pageSize: 25 }), + ); + expect(withSize.get("page_size")).toBe("25"); + expect(withSize.has("page")).toBe(false); + + const backToDefault = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange({ pageIndex: 0, pageSize: 50 }), + ); + expect(backToDefault.has("page_size")).toBe(false); + }); + + it("setSearch writes the term and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.setSearch("prod")); + + expect(url.get("search")).toBe("prod"); + expect(url.has("page")).toBe(false); + expect(result.current.search).toBe("prod"); + expect(result.current.pagination.pageIndex).toBe(0); + }); + + it("setSearch with an empty string removes the key", async () => { + const { result, onUrlUpdate } = renderTableState({ search: "prod" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.setSearch("")); + + expect(url.has("search")).toBe(false); + expect(result.current.search).toBe(""); + }); + + it("onSortingChange writes sort_by and sort_order and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }])); + + expect(url.get("sort_by")).toBe("spend"); + expect(url.get("sort_order")).toBe("asc"); + expect(url.has("page")).toBe(false); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("onSortingChange drops the keys when the sort matches the default or is cleared", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend", sort_order: "asc" }); + + const explicitDefault = await flushUrl(onUrlUpdate, () => + result.current.onSortingChange([{ id: "created_at", desc: true }]), + ); + expect(explicitDefault.has("sort_by")).toBe(false); + expect(explicitDefault.has("sort_order")).toBe(false); + + await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "key_alias", desc: false }])); + const cleared = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([])); + expect(cleared.has("sort_by")).toBe(false); + expect(cleared.has("sort_order")).toBe(false); + expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]); + }); + + it("onSortingChange resolves a function updater against the current sort", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange(flipDirection)); + + expect(url.get("sort_by")).toBe("spend"); + expect(url.get("sort_order")).toBe("asc"); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("onColumnFiltersChange writes trimmed filter_ keys and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: " team-1 " }]), + ); + + expect(url.get("filter_team_id")).toBe("team-1"); + expect(url.has("page")).toBe(false); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("onColumnFiltersChange removes the key for an empty value and for a filter no longer present", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1", filter_user_id: "user-42" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onColumnFiltersChange([{ id: "team_id", value: "" }])); + + expect(url.has("filter_team_id")).toBe(false); + expect(url.has("filter_user_id")).toBe(false); + expect(result.current.columnFilters).toEqual([]); + }); + + it("onColumnFiltersChange ignores a non-string filter value", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: ["team-1", "team-2"] }]), + ); + + expect(url.has("filter_team_id")).toBe(false); + }); + + it("onColumnFiltersChange resolves a function updater against the current filters", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange((previous) => [...previous, { id: "user_id", value: "user-42" }]), + ); + + expect(url.get("filter_team_id")).toBe("team-1"); + expect(url.get("filter_user_id")).toBe("user-42"); + }); + + it("writes prefixed and renamed keys only", async () => { + const { result, onUrlUpdate } = renderTableState( + {}, + { keyPrefix: "audit_", urlKeys: { search: "key_search", filter_team_id: "filter_team" } }, + ); + + await flushUrl(onUrlUpdate, () => result.current.setSearch("prod")); + await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }])); + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: "team-1" }]), + ); + + expect(url.get("audit_key_search")).toBe("prod"); + expect(url.get("audit_sort_by")).toBe("spend"); + expect(url.get("audit_filter_team")).toBe("team-1"); + expect([...url.keys()].filter((key) => !key.startsWith("audit_"))).toEqual([]); + expect(url.has("audit_search")).toBe(false); + expect(url.has("audit_filter_team_id")).toBe(false); + }); +}); + +describe("referential stability", () => { + it("keeps the TanStack state and the page-clamp handler stable across rerenders while the URL is unchanged", () => { + const { result, rerender } = renderTableState({ page: "2", filter_team_id: "team-1", sort_by: "spend" }); + const first = result.current; + + rerender(); + + expect(result.current.sorting).toBe(first.sorting); + expect(result.current.pagination).toBe(first.pagination); + expect(result.current.columnFilters).toBe(first.columnFilters); + expect(result.current.onPaginationChange).toBe(first.onPaginationChange); + }); + + it("hands out new pagination and untouched sorting after a page change", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" }); + const first = result.current; + + await flushUrl(onUrlUpdate, () => result.current.onPaginationChange({ pageIndex: 4, pageSize: 50 })); + + expect(result.current.pagination).not.toBe(first.pagination); + expect(result.current.pagination.pageIndex).toBe(4); + expect(result.current.sorting).toBe(first.sorting); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts new file mode 100644 index 00000000000..1a423663936 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts @@ -0,0 +1,232 @@ +import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { createParser, Nullable, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; +import { useCallback, useMemo } from "react"; + +const SORT_ORDERS = ["asc", "desc"] as const; +type SortOrder = (typeof SORT_ORDERS)[number]; + +const STANDARD_KEYS = ["search", "sort_by", "sort_order", "page", "page_size"] as const; +type StandardKey = (typeof STANDARD_KEYS)[number]; +type FilterStateKey = `filter_${F}`; +type StateKey = StandardKey | FilterStateKey; + +const MAX_PAGE = 100_000; +const DEFAULT_MAX_PAGE_SIZE = 100; + +export interface UrlTableStateOptions { + sortFields: readonly string[]; + defaultSort: { id: string; desc: boolean }; + defaultPageSize: number; + maxPageSize?: number; + filterColumns: readonly F[]; + keyPrefix?: string; + urlKeys?: Partial, string>>; +} + +export interface UrlTableState { + search: string; + setSearch: (value: string) => void; + sorting: SortingState; + onSortingChange: OnChangeFn; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + columnFilters: ColumnFiltersState; + onColumnFiltersChange: OnChangeFn; +} + +const boundedInteger = (min: number, max: number, fallback: number) => + createParser({ + parse: (value: string) => { + const parsed = parseAsInteger.parse(value); + return parsed === null ? null : Math.min(Math.max(parsed, min), max); + }, + serialize: String, + }).withDefault(fallback); + +const optionalString = parseAsString.withDefault(""); +type OptionalStringParser = typeof optionalString; +const sortOrderParser = (fallback: SortOrder) => parseAsStringLiteral(SORT_ORDERS).withDefault(fallback); + +interface StandardValues { + search: string; + sort_by: string; + sort_order: SortOrder; + page: number; + page_size: number; +} +type FilterValues = Record, string>; +type StandardUpdate = Partial>; +type FilterUpdate = Record, string | null> & Pick, "page">; +type SetTableValues = (update: StandardUpdate | FilterUpdate | null) => Promise; + +interface TableQueryState { + values: StandardValues; + filters: FilterValues; + setValues: SetTableValues; +} + +type TableParsers = { + search: OptionalStringParser; + sort_by: OptionalStringParser; + sort_order: ReturnType; + page: ReturnType; + page_size: ReturnType; +} & Record, OptionalStringParser>; + +const useTableQueryStates = ( + parsers: TableParsers, + urlKeys: Record, string>, +): TableQueryState => { + const [state, setState] = useQueryStates(parsers, { urlKeys }); + return useMemo( + () => ({ + values: state as StandardValues, + filters: state as FilterValues, + setValues: setState as SetTableValues, + }), + [state, setState], + ); +}; + +const filterStateKey = (column: F): FilterStateKey => `filter_${column}`; + +const filterParsers = (filterColumns: readonly F[]) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), optionalString])) as Record< + FilterStateKey, + OptionalStringParser + >; + +const resolveUrlKeys = ( + filterColumns: readonly F[], + keyPrefix: string, + renamed: Partial, string>>, +) => { + const stateKeys: readonly StateKey[] = [ + ...STANDARD_KEYS, + ...filterColumns.map((column) => filterStateKey(column)), + ]; + return Object.fromEntries(stateKeys.map((key) => [key, `${keyPrefix}${renamed[key] ?? key}`])) as Record< + StateKey, + string + >; +}; + +const filterValue = (filters: ColumnFiltersState, column: string): string | null => { + const value = filters.find((filter) => filter.id === column)?.value; + return (typeof value === "string" ? value.trim() : "") || null; +}; + +const filterUpdates = (filterColumns: readonly F[], filters: ColumnFiltersState) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), filterValue(filters, column)])) as Record< + FilterStateKey, + string | null + >; + +const toSortOrder = (active: SortingState[number]): SortOrder => (active.desc ? "desc" : "asc"); + +export function useUrlTableState(options: UrlTableStateOptions): UrlTableState { + const { + sortFields, + defaultSort, + defaultPageSize, + maxPageSize = DEFAULT_MAX_PAGE_SIZE, + filterColumns, + keyPrefix = "", + urlKeys: renamedKeys, + } = options; + const defaultSortId = defaultSort.id; + const defaultSortOrder: SortOrder = defaultSort.desc ? "desc" : "asc"; + + const parsers = useMemo>( + () => ({ + search: optionalString, + sort_by: parseAsString.withDefault(defaultSortId), + sort_order: sortOrderParser(defaultSortOrder), + page: boundedInteger(1, MAX_PAGE, 1), + page_size: boundedInteger(1, maxPageSize, defaultPageSize), + ...filterParsers(filterColumns), + }), + [defaultSortId, defaultSortOrder, defaultPageSize, maxPageSize, filterColumns], + ); + const urlKeys = useMemo( + () => resolveUrlKeys(filterColumns, keyPrefix, renamedKeys ?? {}), + [filterColumns, keyPrefix, renamedKeys], + ); + const { values, filters, setValues } = useTableQueryStates(parsers, urlKeys); + + const sortBy = sortFields.includes(values.sort_by) ? values.sort_by : defaultSortId; + const sortDesc = values.sort_order === "desc"; + const sorting = useMemo(() => [{ id: sortBy, desc: sortDesc }], [sortBy, sortDesc]); + + const pagination = useMemo( + () => ({ pageIndex: values.page - 1, pageSize: values.page_size }), + [values.page, values.page_size], + ); + + const columnFilters = useMemo( + () => + filterColumns.flatMap((column) => { + const value = filters[filterStateKey(column)].trim(); + return value ? [{ id: column, value }] : []; + }), + [filterColumns, filters], + ); + + const setSearch = useCallback( + (value: string) => { + void setValues({ search: value || null, page: null }); + }, + [setValues], + ); + + const onSortingChange = useCallback>( + (updaterOrValue) => { + const active = functionalUpdate(updaterOrValue, sorting)[0]; + void setValues({ + sort_by: active?.id ?? null, + sort_order: active ? toSortOrder(active) : null, + page: null, + }); + }, + [setValues, sorting], + ); + + const onPaginationChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, pagination); + void setValues({ page: next.pageIndex + 1, page_size: next.pageSize }); + }, + [pagination, setValues], + ); + + const onColumnFiltersChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, columnFilters); + void setValues({ ...filterUpdates(filterColumns, next), page: null }); + }, + [columnFilters, filterColumns, setValues], + ); + + return useMemo( + () => ({ + search: values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + }), + [ + values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + ], + ); +} diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx new file mode 100644 index 00000000000..427d454ee2d --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx @@ -0,0 +1,100 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { NuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import type { ReactNode } from "react"; +import { describe, expect, it, vi } from "vitest"; +import { useUrlTab } from "./useUrlTab"; + +const TABS = ["chat", "compare", "compliance"] as const; +type Tab = (typeof TABS)[number]; + +interface RenderArgs { + searchParams?: string; + onUrlUpdate?: OnUrlUpdateFunction; + key?: string; +} + +const initialProps: { values: readonly Tab[] } = { values: TABS }; + +const renderUrlTab = ({ searchParams, onUrlUpdate, key }: RenderArgs = {}) => + renderHook(({ values }: { values: readonly Tab[] }) => useUrlTab(values, "chat", key), { + initialProps, + wrapper: ({ children }: { children: ReactNode }) => ( + + {children} + + ), + }); + +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + onUrlUpdate.mock.calls.at(-1)?.[0]; + +describe("useUrlTab", () => { + it("reads the active tab from the URL", () => { + const { result } = renderUrlTab({ searchParams: "?tab=compare" }); + + expect(result.current[0]).toBe("compare"); + }); + + it("resolves a URL value outside the allowed tabs to the fallback and drops it from the URL", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ searchParams: "?tab=settings&other=1", onUrlUpdate }); + + expect(result.current[0]).toBe("chat"); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("other")).toBe("1"); + }); + + it("leaves a URL that names an allowed tab untouched", async () => { + const onUrlUpdate = vi.fn(); + renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate }); + + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(onUrlUpdate).not.toHaveBeenCalled(); + }); + + it("reads from the caller's key instead of the default one", () => { + const { result } = renderUrlTab({ searchParams: "?view=compliance&tab=compare", key: "view" }); + + expect(result.current[0]).toBe("compliance"); + }); + + it("writes ?tab= with history replace when a tab is selected", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ onUrlUpdate }); + + act(() => result.current[1]("compare")); + + await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compare")); + expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace"); + expect(result.current[0]).toBe("compare"); + }); + + it("removes the param when the fallback tab is selected", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate }); + + act(() => result.current[1]("chat")); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + expect(result.current[0]).toBe("chat"); + }); + + it("falls back and clears the param when the current tab is no longer among the allowed values", async () => { + const onUrlUpdate = vi.fn(); + const { result, rerender } = renderUrlTab({ searchParams: "?tab=compliance", onUrlUpdate }); + expect(result.current[0]).toBe("compliance"); + + rerender({ values: ["chat", "compare"] }); + + expect(result.current[0]).toBe("chat"); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.ts b/ui/litellm-dashboard/src/hooks/useUrlTab.ts new file mode 100644 index 00000000000..2f3d705c610 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.ts @@ -0,0 +1,12 @@ +import { parseAsString, useQueryState } from "nuqs"; +import { useCallback, useEffect } from "react"; + +export function useUrlTab(values: readonly T[], fallback: T, key = "tab"): [T, (tab: T) => void] { + const [urlTab, setUrlTab] = useQueryState(key, parseAsString.withDefault(fallback)); + const tab = values.find((value) => value === urlTab) ?? fallback; + useEffect(() => { + if (urlTab !== tab) void setUrlTab(null); + }, [urlTab, tab, setUrlTab]); + const setTab = useCallback((next: T) => void setUrlTab(next), [setUrlTab]); + return [tab, setTab]; +} diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts b/ui/litellm-dashboard/src/utils/tabRoutes.test.ts deleted file mode 100644 index 402be55c33a..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts +++ /dev/null @@ -1,47 +0,0 @@ -/* @vitest-environment jsdom */ -import { describe, expect, it, vi } from "vitest"; - -vi.mock("@/components/networking", () => ({ serverRootPath: "" })); - -import { createTabRoutes } from "./tabRoutes"; - -const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const); - -describe("createTabRoutes.slugFromPathname", () => { - it("returns empty string for the base path with or without a trailing slash", () => { - expect(routes.slugFromPathname("/logs")).toBe(""); - expect(routes.slugFromPathname("/logs/")).toBe(""); - }); - - it("extracts the tab slug from dev and proxy-mounted (/ui) paths", () => { - expect(routes.slugFromPathname("/logs/audit")).toBe("audit"); - expect(routes.slugFromPathname("/ui/logs/deleted-teams/")).toBe("deleted-teams"); - }); - - it("returns the raw segment for an unknown tab so the caller can redirect to base", () => { - expect(routes.slugFromPathname("/ui/logs/bogus")).toBe("bogus"); - }); - - it("returns empty string when the base segment is not in the path", () => { - expect(routes.slugFromPathname("/teams")).toBe(""); - }); -}); - -describe("createTabRoutes.tabHref", () => { - it("builds the trailing-slash base href for the empty slug", () => { - expect(routes.tabHref("")).toBe("/ui/logs/"); - }); - - it("builds a trailing-slash href for every tab slug (required by static export)", () => { - for (const slug of routes.slugs) { - expect(routes.tabHref(slug)).toBe(`/ui/logs/${slug}/`); - } - }); -}); - -describe("createTabRoutes metadata", () => { - it("preserves the base segment and slug tuple", () => { - expect(routes.baseSegment).toBe("logs"); - expect(routes.slugs).toEqual(["audit", "deleted-keys", "deleted-teams"]); - }); -}); diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.ts b/ui/litellm-dashboard/src/utils/tabRoutes.ts deleted file mode 100644 index 4af2b983cba..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { uiHref } from "@/utils/uiHref"; - -export interface TabRoutes { - baseSegment: string; - slugs: readonly Slug[]; - tabHref: (slug: string) => string; - slugFromPathname: (pathname: string) => string; -} - -export function createTabRoutes(baseSegment: string, slugs: readonly Slug[]): TabRoutes { - const tabHref = (slug: string): string => { - const base = uiHref(baseSegment); - return slug ? `${base}/${slug}/` : `${base}/`; - }; - - const slugFromPathname = (pathname: string): string => { - const parts = pathname.split("/").filter(Boolean); - const idx = parts.indexOf(baseSegment); - if (idx === -1) { - return ""; - } - return parts[idx + 1] ?? ""; - }; - - return { baseSegment, slugs, tabHref, slugFromPathname }; -} From 4d30bbce453a06e850cbf6007a42dd7e00803380 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 09:34:05 -0700 Subject: [PATCH 71/89] fix(ui): keep controlled client-side table pages across data reloads Controlled client-mode DataTables no longer let TanStack reset the page index when rows change, since the owner of the pagination state decides the page. Once rows settle, a page past the end snaps back to the last page, matching server mode --- .../shared/DataTable/DataTable.test.tsx | 73 +++++++++++++++++++ .../components/shared/DataTable/DataTable.tsx | 34 ++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 336fc43d695..a7e4befa6ac 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -357,6 +357,79 @@ describe("DataTable pagination", () => { expect(onChange).not.toHaveBeenCalled(); }); + type ClientPageHarnessProps = { + data: Person[]; + isLoading?: boolean; + initialPageIndex: number; + onChange: (next: PaginationState) => void; + }; + + function ClientPageHarness({ data, isLoading = false, initialPageIndex, onChange }: ClientPageHarnessProps) { + const [pagination, setPagination] = useState({ pageIndex: initialPageIndex, pageSize: 2 }); + const handleChange: OnChangeFn = (updater) => { + const next = typeof updater === "function" ? updater(pagination) : updater; + onChange(next); + setPagination(next); + }; + return ( + + ); + } + + it("client mode keeps a controlled page when rows arrive after loading and when they are refetched", async () => { + const onChange = vi.fn(); + const { rerender } = render(); + + rerender(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(names()).toEqual(["P2", "P3"]); + + rerender(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(names()).toEqual(["P2", "P3"]); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("client mode snaps a controlled page past the end back to the last page", async () => { + const onChange = vi.fn(); + render(); + + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 2, pageSize: 2 })); + expect(onChange).toHaveBeenCalledTimes(1); + expect(names()).toEqual(["P4"]); + }); + + it("client mode leaves a controlled page alone while there are no rows to page through", async () => { + const onChange = vi.fn(); + render(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("client mode without a controlled page still returns to the first page when the rows change", async () => { + const user = userEvent.setup(); + const { rerender } = render( + , + ); + + await user.click(screen.getByTestId("pagination-next")); + expect(names()).toEqual(["P2", "P3"]); + + rerender( + , + ); + await waitFor(() => expect(names()).toEqual(["P0", "P1"])); + }); + it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => { const onChange = vi.fn(); const { rerender } = render(); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index e0f57ae1052..340f8d4f44f 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -425,7 +425,7 @@ function useControllable( return { value: internal, onChange: setInternal }; } -function useServerPageClamp( +function usePageClamp( active: boolean, rowCount: number | undefined, pagination: { value: PaginationState; onChange: OnChangeFn }, @@ -484,7 +484,6 @@ function useDataTableInstance( pageIndex: 0, pageSize: pageSizeOptions[0] ?? 25, }); - useServerPageClamp(paginationMode === "server" && !isLoading && !isError, rowCount, paginationState); const filterState = useControllable( columnFilters, onColumnFiltersChange, @@ -536,9 +535,38 @@ function useDataTableInstance( ...(getRowId !== undefined ? { getRowId } : {}), ...(enableRowSelection !== undefined ? { enableRowSelection } : {}), ...(paginationMode === "server" && rowCount !== undefined ? { rowCount } : {}), + autoResetPageIndex: pagination === undefined && paginationMode !== "server", }; - return useReactTable(tableOptions); + const table = useReactTable(tableOptions); + const clampOptions: SettledPageClampOptions = { + paginationMode, + controlled: pagination !== undefined, + settled: !isLoading && !isError, + rowCount, + pagination: paginationState, + }; + useSettledPageClamp(table, clampOptions); + return table; +} + +type SettledPageClampOptions = { + paginationMode: PaginationMode; + controlled: boolean; + settled: boolean; + rowCount: number | undefined; + pagination: { value: PaginationState; onChange: OnChangeFn }; +}; + +function useSettledPageClamp(table: Table, options: SettledPageClampOptions): void { + const { paginationMode, controlled, settled, rowCount, pagination } = options; + const clientRowCount = paginationMode === "client" ? table.getPrePaginationRowModel().rows.length : 0; + const clientPageIsClampable = paginationMode === "client" && controlled && clientRowCount > 0; + usePageClamp( + settled && (paginationMode === "server" || clientPageIsClampable), + paginationMode === "server" ? rowCount : clientRowCount, + pagination, + ); } export function DataTable(props: DataTableProps) { From adc937c4931021c68864ed28014d51b851ecaade Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 12:50:08 -0700 Subject: [PATCH 72/89] fix(ui): read persisted column visibility from storage instead of a mounted copy usePersistedColumnVisibility kept a useState copy seeded once at mount, so a later tableId or defaults change showed the old table's columns and saved them under the new key. It now reads localStorage through useSyncExternalStore, keeping only writes that storage refused in memory, so the hook has no copy to go stale. Stored choices are layered over the defaults on every read, and changes saved in another tab show up. --- .../usePersistedColumnVisibility.test.tsx | 74 ++++++++++++++++++- .../DataTable/usePersistedColumnVisibility.ts | 54 +++++++++++--- 2 files changed, 112 insertions(+), 16 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx index 8d2e0f61f54..fbfb71e5bc2 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx @@ -1,3 +1,4 @@ +import type { VisibilityState } from "@tanstack/react-table"; import { act, renderHook } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; @@ -10,6 +11,9 @@ const stored = (tableId: string): unknown => { return raw === null ? null : JSON.parse(raw); }; +const showEveryColumn = (previous: VisibilityState): VisibilityState => + Object.fromEntries(Object.keys(previous).map((column) => [column, true])); + describe("usePersistedColumnVisibility", () => { beforeEach(() => { localStorage.clear(); @@ -55,6 +59,15 @@ describe("usePersistedColumnVisibility", () => { expect(stored("keys")).toEqual({ email: false, name: false }); }); + it("hands a function updater the default-hidden columns, so showing every column sticks", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + + act(() => result.current.onColumnVisibilityChange(showEveryColumn)); + + expect(result.current.columnVisibility).toEqual({ spend: true }); + expect(stored("keys")).toEqual({ spend: true }); + }); + it.each([ ["truncated JSON", '{"email":fal'], ["a JSON scalar", "42"], @@ -80,19 +93,72 @@ describe("usePersistedColumnVisibility", () => { expect(stored("teams")).toBeNull(); }); + it("reads and writes the new table's columns after the tableId changes", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + localStorage.setItem(keyFor("teams"), JSON.stringify({ spend: false })); + const { result, rerender } = renderHook(({ tableId }) => usePersistedColumnVisibility(tableId), { + initialProps: { tableId: "keys" }, + }); + + rerender({ tableId: "teams" }); + expect(result.current.columnVisibility).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false }))); + expect(stored("teams")).toEqual({ spend: false, name: false }); + expect(stored("keys")).toEqual({ email: false }); + }); + + it("applies new defaults passed after mount", () => { + const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), { + initialProps: { defaults: { spend: false } }, + }); + + rerender({ defaults: { name: false } }); + + expect(result.current.columnVisibility).toEqual({ name: false }); + }); + + it("shows a change another tab saved for the same table", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + window.dispatchEvent(new StorageEvent("storage", { key: keyFor("keys") })); + }); + + expect(result.current.columnVisibility).toEqual({ email: false }); + }); + + it("keeps a toggle that storage refused, and saves the next one once storage accepts it", () => { + localStorage.setItem(keyFor("full"), JSON.stringify({ spend: false })); + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("full")); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + expect(result.current.columnVisibility).toEqual({ email: false }); + expect(stored("full")).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange({ name: false })); + expect(result.current.columnVisibility).toEqual({ name: false }); + expect(stored("full")).toEqual({ name: false }); + }); + it("returns the defaults without throwing when storage is unavailable", () => { vi.spyOn(console, "warn").mockImplementation(() => {}); vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { throw new Error("SecurityError"); }); vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { - throw new Error("QuotaExceededError"); + throw new Error("SecurityError"); }); - const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + const { result } = renderHook(() => usePersistedColumnVisibility("blocked", { spend: false })); expect(result.current.columnVisibility).toEqual({ spend: false }); - act(() => result.current.onColumnVisibilityChange({ email: false })); - expect(result.current.columnVisibility).toEqual({ email: false }); + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, email: false }))); + expect(result.current.columnVisibility).toEqual({ spend: false, email: false }); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts index b56ae13d63b..7fffc50b381 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -1,16 +1,46 @@ import type { OnChangeFn, VisibilityState } from "@tanstack/react-table"; -import { useCallback, useState } from "react"; +import { useCallback, useMemo, useSyncExternalStore } from "react"; -import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { + LOCAL_STORAGE_EVENT, + emitLocalStorageChange, + getLocalStorageItem, + setLocalStorageItem, +} from "@/utils/localStorageUtils"; const STORAGE_KEY_PREFIX = "litellm_table_columns_"; const EMPTY_VISIBILITY: VisibilityState = {}; +const unsavedWrites = new Map(); + function storageKey(tableId: string): string { return `${STORAGE_KEY_PREFIX}${tableId}`; } +function subscribe(onChange: () => void): () => void { + window.addEventListener("storage", onChange); + window.addEventListener(LOCAL_STORAGE_EVENT, onChange); + return () => { + window.removeEventListener("storage", onChange); + window.removeEventListener(LOCAL_STORAGE_EVENT, onChange); + }; +} + +function readRaw(key: string): string | null { + return unsavedWrites.get(key) ?? getLocalStorageItem(key); +} + +function writeRaw(key: string, raw: string): void { + setLocalStorageItem(key, raw); + if (getLocalStorageItem(key) === raw) { + unsavedWrites.delete(key); + } else { + unsavedWrites.set(key, raw); + } + emitLocalStorageChange(key); +} + function isVisibilityState(value: unknown): value is VisibilityState { if (typeof value !== "object" || value === null || Array.isArray(value)) { return false; @@ -18,8 +48,7 @@ function isVisibilityState(value: unknown): value is VisibilityState { return Object.values(value).every((visible) => typeof visible === "boolean"); } -function readStoredVisibility(tableId: string, defaults: VisibilityState): VisibilityState { - const raw = getLocalStorageItem(storageKey(tableId)); +function parseVisibility(raw: string | null, defaults: VisibilityState): VisibilityState { if (raw === null) { return defaults; } @@ -35,19 +64,20 @@ export function usePersistedColumnVisibility( tableId: string, defaults: VisibilityState = EMPTY_VISIBILITY, ): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn } { - const [columnVisibility, setColumnVisibility] = useState(() => - readStoredVisibility(tableId, defaults), + const key = storageKey(tableId); + const raw = useSyncExternalStore( + subscribe, + () => readRaw(key), + () => null, ); + const columnVisibility = useMemo(() => parseVisibility(raw, defaults), [raw, defaults]); const onColumnVisibilityChange = useCallback>( (updater) => { - setColumnVisibility((previous) => { - const next = typeof updater === "function" ? updater(previous) : updater; - setLocalStorageItem(storageKey(tableId), JSON.stringify(next)); - return next; - }); + const next = typeof updater === "function" ? updater(parseVisibility(readRaw(key), defaults)) : updater; + writeRaw(key, JSON.stringify(next)); }, - [tableId], + [key, defaults], ); return { columnVisibility, onColumnVisibilityChange }; From d29753c52172b619a2ad7702b34fe1e09b45a6f5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:43:21 -0700 Subject: [PATCH 73/89] fix(ui): let another tab's column save replace a toggle this tab could not save A column toggle that localStorage refused was kept in memory and read ahead of storage, so a later save from another tab stayed hidden until this tab saved again. A storage event now drops the in-memory copy for its key, or all of them when another tab clears storage --- .../usePersistedColumnVisibility.test.tsx | 32 ++++++++++++++++++- .../DataTable/usePersistedColumnVisibility.ts | 16 ++++++++-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx index fbfb71e5bc2..35bba906afc 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx @@ -109,8 +109,9 @@ describe("usePersistedColumnVisibility", () => { }); it("applies new defaults passed after mount", () => { + const initialProps: { defaults: VisibilityState } = { defaults: { spend: false } }; const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), { - initialProps: { defaults: { spend: false } }, + initialProps, }); rerender({ defaults: { name: false } }); @@ -146,6 +147,35 @@ describe("usePersistedColumnVisibility", () => { expect(stored("full")).toEqual({ name: false }); }); + it("shows another tab's save over a toggle this tab could not save", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("shadowed")); + act(() => result.current.onColumnVisibilityChange({ email: false })); + + act(() => { + localStorage.setItem(keyFor("shadowed"), JSON.stringify({ name: false })); + window.dispatchEvent(new StorageEvent("storage", { key: keyFor("shadowed") })); + }); + + expect(result.current.columnVisibility).toEqual({ name: false }); + }); + + it("drops a toggle this tab could not save once another tab clears storage", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("cleared", { spend: false })); + act(() => result.current.onColumnVisibilityChange({ email: false })); + + act(() => window.dispatchEvent(new StorageEvent("storage", { key: null }))); + + expect(result.current.columnVisibility).toEqual({ spend: false }); + }); + it("returns the defaults without throwing when storage is unavailable", () => { vi.spyOn(console, "warn").mockImplementation(() => {}); vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts index 7fffc50b381..3cbf5c2a000 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -18,11 +18,23 @@ function storageKey(tableId: string): string { return `${STORAGE_KEY_PREFIX}${tableId}`; } +function forgetUnsavedWrite(event: StorageEvent): void { + if (event.key === null) { + unsavedWrites.clear(); + return; + } + unsavedWrites.delete(event.key); +} + function subscribe(onChange: () => void): () => void { - window.addEventListener("storage", onChange); + const onStorage = (event: StorageEvent): void => { + forgetUnsavedWrite(event); + onChange(); + }; + window.addEventListener("storage", onStorage); window.addEventListener(LOCAL_STORAGE_EVENT, onChange); return () => { - window.removeEventListener("storage", onChange); + window.removeEventListener("storage", onStorage); window.removeEventListener(LOCAL_STORAGE_EVENT, onChange); }; } From 36eb9cdb356e0ad02124e7ae149324def11b4669 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:56:27 -0700 Subject: [PATCH 74/89] test(proxy): drop the route-list membership test that the PATCH gate tests already cover --- .../proxy/auth/test_route_checks.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 3211e85ff97..72c7011e6f5 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2892,7 +2892,7 @@ def test_team_update_gate_allows_org_admin_with_resolved_org(): ) -def test_team_update_gate_admits_internal_user_without_org_context(): +def test_team_update_gate_admits_internal_user_without_org_context(): # test-quality-ok: the gate's only success signal is not raising; the handler's team-admin 403s are pinned in test_team_endpoints """/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated caller and update_team resolves proxy, org or team admin itself, then filters team admins through the team_admin_editable_team_fields setting. Before that the gate 401'd every @@ -2997,23 +2997,6 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_has_same_reach_as_team_update(): - """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but - NOT by regular internal users or the role-agnostic self_managed_routes — the - latter would open /team/new (the collision footgun) to any authenticated user.""" - from litellm.proxy._types import LiteLLMRoutes - - assert RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value - ) - assert not RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.internal_user_routes.value - ) - assert not RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.self_managed_routes.value - ) - - def _patch_team_request() -> MagicMock: request = MagicMock(spec=Request) request.method = "PATCH" From afb28540bbd3868dcebd83e8e7c7347c7611abfa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 16:01:32 -0700 Subject: [PATCH 75/89] fix(e2e): keep the CLI determinism test out of the in-cluster suite It drives the real CLI for several seconds. The edge stamps every upstream call with PYTEST_CURRENT_TEST, a process-global that names whichever test the worker is in when the call arrives rather than the one that made it, so a test that holds a worker that long collects other tests' in-flight calls. Build 234's key report credits this test with 20 Bedrock and 7 Anthropic misses, and it makes no provider call at all. Those misattributed calls take the wrong test id into the cache key and write recordings under it, so the test was polluting the shared corpus it exists to protect. Deselected unless E2E_CLI_DETERMINISM is set, the same opt-in shape the managed-files, prompt-caching and redis-chaos markers already use. The attribution bug itself is older than this branch and is reported, not fixed here. --- .../_driver_unit_tests/test_request_determinism.py | 2 ++ tests/e2e/conftest.py | 7 +++++++ tests/e2e/e2e_config.py | 1 + tests/e2e/pytest.ini | 1 + 4 files changed, 11 insertions(+) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 5046f35c73b..b7d330b7da6 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -36,6 +36,8 @@ import pytest from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter +pytestmark = pytest.mark.cli_determinism + _STUB_REPLY = { "id": "msg_stub", "type": "message", diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 430e16525d5..ac4cfb71407 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -23,6 +23,7 @@ from typing import Final import pytest import requests from e2e_config import ( + CLI_DETERMINISM_OPT_IN_ENV, CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, @@ -53,6 +54,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, } ) @@ -120,6 +122,11 @@ def pytest_configure(config: pytest.Config) -> None: "prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including " "prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set", ) + config.addinivalue_line( + "markers", + "cli_determinism: drives the real claude CLI for several seconds, which widens the window in which " + "another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set", + ) config.addinivalue_line( "markers", "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..82ddb09f7f5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,7 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..f6d23a3ec12 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -10,4 +10,5 @@ markers = weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set + cli_determinism: drives the real claude CLI for several seconds, which widens the window in which another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set From 02d9aae8c87fd2948a37c47815995a16465e45ce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:01:52 -0700 Subject: [PATCH 76/89] fix(proxy): stop forwarding LiteLLM credential headers on Bedrock agent-runtime passthrough The agent-runtime branch of /bedrock/{endpoint} (agents, knowledgebases, flows, retrieveAndGenerate, rerank, generateQuery, optimize-prompt) forwarded every caller header to AWS next to the SigV4 signature, so a LiteLLM key presented in x-api-key or x-litellm-api-key reached bedrock-agent-runtime verbatim. Build the upstream header set explicitly: drop LiteLLM credential headers by name and any authenticated secret by value, keep the remaining caller headers, and let the signed headers win on collisions. --- .../llm_passthrough_endpoints.py | 17 ++- .../test_llm_pass_through_endpoints.py | 140 ++++++++++++++++++ 2 files changed, 155 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0fe9d1cc626..b9b8cb3a22b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1180,9 +1180,8 @@ async def bedrock_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=str(prepped.url), - custom_headers=prepped.headers, + custom_headers=_upstream_headers_for_bedrock_agent_runtime_route(request, user_api_key_dict, prepped.headers), is_streaming_request=is_streaming_request, - _forward_headers=True, ) # dynamically construct pass-through endpoint based on incoming path setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data) # SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps @@ -2001,6 +2000,9 @@ _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-a _HEADERS_NEVER_FORWARDED_TO_ANTHROPIC: Final = frozenset({"content-length", "host", "accept-encoding"}) | ( SpecialHeaders.litellm_credential_header_names() - _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS ) +_HEADERS_NEVER_FORWARDED_TO_BEDROCK: Final = ( + frozenset({"content-length", "host", "accept-encoding"}) | SpecialHeaders.litellm_credential_header_names() +) _MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key" @@ -2099,6 +2101,17 @@ def _upstream_headers_for_anthropic_route( return MappingProxyType({**caller_headers, **(proxy_auth_header or {})}) +def _upstream_headers_for_bedrock_agent_runtime_route( + request: Request, user_api_key_dict: UserAPIKeyAuth, signed_headers: Mapping[str, object] +) -> Mapping[str, object]: + caller_headers: Final = _caller_headers_without_litellm_secrets( + request, + user_api_key_dict, + _HEADERS_NEVER_FORWARDED_TO_BEDROCK | frozenset(name.lower() for name in signed_headers), + ) + return MappingProxyType({**caller_headers, **signed_headers}) + + async def _prepare_vertex_auth_headers( request: Request, vertex_credentials: VertexPassThroughCredentials | None, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index e0785b002b2..7f044486e14 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1983,6 +1983,146 @@ class TestBedrockAgentRuntimePassthroughToggle: create_route.assert_called_once() +class TestBedrockAgentRuntimePassthroughVirtualKeyLeak: + """Regression for LIT-7912: the agent-runtime branch of ``/bedrock/{endpoint}`` forwarded every caller header, + so a LiteLLM key presented in ``x-api-key`` or ``x-litellm-api-key`` rode to AWS next to the SigV4 signature.""" + + VKEY: Final = "sk-litellm-victim-key" + MASTER_KEY: Final = "sk-master-1234" + ENDPOINT: Final = "knowledgebases/KB1234567/retrieve" + AMBIENT_AWS_ENV: Final = ( + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_SESSION_TOKEN", + "AWS_SESSION_NAME", + "AWS_PROFILE_NAME", + "AWS_ROLE_NAME", + "AWS_WEB_IDENTITY_TOKEN", + "AWS_STS_ENDPOINT", + "AWS_EXTERNAL_ID", + ) + + async def _upstream_headers(self, monkeypatch, headers: list[tuple[bytes, bytes]]) -> dict: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", self.MASTER_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + for ambient in self.AMBIENT_AWS_ENV: + monkeypatch.delenv(ambient, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + caller: Final = UserAPIKeyAuth(api_key=self.VKEY) + + async def receive(): + return {"type": "http.request", "body": b'{"retrievalQuery": {"text": "hi"}}', "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/{self.ENDPOINT}", + "headers": headers, + "query_string": b"", + }, + receive=receive, + ) + captured: dict = {} + + def fake_create_pass_through_route(**kwargs): + captured.update(kwargs) + return AsyncMock(return_value={"status": "success"}) + + module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + with ( + patch(f"{module}.create_request_copy", Mock()), + patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), + ): + await bedrock_proxy_route( + endpoint=self.ENDPOINT, + request=request, + fastapi_response=Response(), + user_api_key_dict=caller, + ) + return HttpPassThroughEndpointHelpers.forward_headers_from_request( + request_headers=dict(request.headers), + headers=dict(captured["custom_headers"] or {}), + forward_headers=captured.get("_forward_headers", False), + ) + + @staticmethod + def _blob(upstream: dict) -> str: + return " ".join(f"{name}:{value}" for name, value in upstream.items()) + + @staticmethod + def _names_matching(upstream: dict, lowercase_name: str) -> list[str]: + return [name for name in upstream if name.lower() == lowercase_name] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "header_name", ["x-api-key", "x-litellm-api-key", "api-key", "x-goog-api-key", "ocp-apim-subscription-key"] + ) + async def test_virtual_key_in_a_credential_header_never_reaches_aws(self, monkeypatch, header_name: str): + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (header_name.encode(), self.VKEY.encode()), + (b"content-type", b"application/json"), + (b"x-request-id", b"trace-1"), + ], + ) + + assert self.VKEY not in self._blob(upstream) + assert self._names_matching(upstream, header_name) == [] + assert upstream["x-request-id"] == "trace-1", "a benign caller header still reaches AWS" + assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256") + assert self._names_matching(upstream, "content-type") == ["Content-Type"], "the signed header is the only one" + + @pytest.mark.asyncio + async def test_credential_headers_are_dropped_by_name_even_when_they_carry_someone_elses_key(self, monkeypatch): + other_key: Final = "sk-other-tenant-key" + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (b"x-api-key", other_key.encode()), + (b"x-litellm-api-key", other_key.encode()), + (b"x-request-id", b"trace-3"), + ], + ) + + assert other_key not in self._blob(upstream) + assert self._names_matching(upstream, "x-api-key") == [] + assert self._names_matching(upstream, "x-litellm-api-key") == [] + assert upstream["x-request-id"] == "trace-3" + + @pytest.mark.asyncio + async def test_virtual_key_in_authorization_bearer_is_replaced_by_the_sigv4_signature(self, monkeypatch): + upstream: Final = await self._upstream_headers( + monkeypatch, + [(b"authorization", f"Bearer {self.VKEY}".encode()), (b"content-type", b"application/json")], + ) + + assert self.VKEY not in self._blob(upstream) + assert self._names_matching(upstream, "authorization") == ["Authorization"] + assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256") + + @pytest.mark.asyncio + async def test_authenticated_secrets_in_any_other_header_never_reach_aws(self, monkeypatch): + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (b"x-api-key", self.VKEY.encode()), + (b"x-forwarded-key", self.VKEY.encode()), + (b"x-operator-token", self.MASTER_KEY.encode()), + (b"x-request-id", b"trace-2"), + ], + ) + + assert self.VKEY not in self._blob(upstream) and self.MASTER_KEY not in self._blob(upstream) + assert self._names_matching(upstream, "x-forwarded-key") == [] + assert self._names_matching(upstream, "x-operator-token") == [] + assert upstream["x-request-id"] == "trace-2" + + class TestLLMPassthroughFactoryProxyRoute: @pytest.mark.asyncio async def test_llm_passthrough_factory_proxy_route_success(self): From 81ae5caa7e65e05be6a07f53fa46b394e8003462 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:11:52 -0700 Subject: [PATCH 77/89] fix(proxy): pass only a team admin's changed fields on to the team update --- .../team_admin_field_permissions.py | 21 ++++++++--- .../management_endpoints/team_endpoints.py | 4 +-- .../test_team_admin_field_permissions.py | 33 +++++++++++------ .../test_team_endpoints.py | 36 +++++++++++++++++++ 4 files changed, 77 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 6e836ad440c..4248501551f 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -35,6 +35,7 @@ _SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields" @dataclass(frozen=True, slots=True) class TeamAdminEditAllowed: + request: UpdateTeamRequest kind: Literal["allowed"] = "allowed" @@ -143,6 +144,15 @@ def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable return column_changes | _metadata_changes(data, submitted, existing) +def _only_changes(data: UpdateTeamRequest, changed: frozenset[str]) -> UpdateTeamRequest: + """The request without the values it resends unchanged, which would otherwise still trigger derived writes + such as a resent budget_duration pushing budget_reset_at back.""" + sent: Final = frozenset(data.model_fields_set) + via_metadata: Final = frozenset({"metadata"}) if changed - sent else frozenset() + kept: Final = frozenset({"team_id"}) | (changed & sent) | via_metadata + return UpdateTeamRequest.model_validate(data.model_dump(include=MappingProxyType({field: True for field in kept}))) + + def team_admin_edit_verdict( data: UpdateTeamRequest, existing: LiteLLM_TeamTable, @@ -150,16 +160,17 @@ def team_admin_edit_verdict( ) -> TeamAdminEditVerdict: if not permitted: return TeamAdminEditingDisabled() - blocked: Final = sorted(changed_team_fields(data, existing) - permitted) + changed: Final = changed_team_fields(data, existing) + blocked: Final = sorted(changed - permitted) if blocked: return TeamAdminFieldNotPermitted(field=blocked[0]) - return TeamAdminEditAllowed() + return TeamAdminEditAllowed(request=_only_changes(data, changed)) -def raise_for_team_admin_edit_verdict(verdict: TeamAdminEditVerdict) -> None: +def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest: match verdict: - case TeamAdminEditAllowed(): - return + case TeamAdminEditAllowed(request=request): + return request case TeamAdminEditingDisabled(): raise HTTPException( status_code=403, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 72367b2bced..b2dc3551ced 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -140,9 +140,9 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( ) from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, - raise_for_team_admin_edit_verdict, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_request_or_raise, ) from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, @@ -2218,7 +2218,7 @@ async def update_team( if access_role is None: _raise_team_access_denied() if access_role == "team_admin": - raise_for_team_admin_edit_verdict( + data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below team_admin_edit_verdict( data=data, existing=existing_team, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py index 91479921c61..5b31089f91e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -7,9 +7,9 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( TeamAdminEditingDisabled, TeamAdminFieldNotPermitted, changed_team_fields, - raise_for_team_admin_edit_verdict, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_request_or_raise, ) _SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"}) @@ -96,10 +96,22 @@ class TestTeamAdminEditVerdict: verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset()) assert verdict == TeamAdminEditingDisabled() - def test_changes_within_permitted_fields_are_allowed(self): - data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") - verdict = team_admin_edit_verdict(data, _team(team_alias="alpha"), frozenset({"tpm_limit"})) - assert verdict == TeamAdminEditAllowed() + def test_allowed_request_keeps_only_the_changed_fields(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha", budget_duration="30d") + existing = _team(team_alias="alpha", budget_duration="30d") + verdict = team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"})) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6} + + def test_permitted_field_changed_inside_metadata_keeps_the_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}, team_alias="alpha") + existing = _team(team_alias="alpha", metadata={"guardrails": ["a"]}) + verdict = team_admin_edit_verdict(data, existing, frozenset({"guardrails"})) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == { + "team_id": "team-1", + "metadata": {"guardrails": ["b"]}, + } def test_first_blocked_field_in_sorted_order_is_reported(self): data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True) @@ -107,19 +119,20 @@ class TestTeamAdminEditVerdict: assert verdict == TeamAdminFieldNotPermitted(field="blocked") -class TestRaiseForTeamAdminEditVerdict: - def test_allowed_does_not_raise(self): - assert raise_for_team_admin_edit_verdict(TeamAdminEditAllowed()) is None +class TestTeamAdminRequestOrRaise: + def test_allowed_hands_back_its_request(self): + request = UpdateTeamRequest(team_id="team-1", tpm_limit=6) + assert team_admin_request_or_raise(TeamAdminEditAllowed(request=request)) is request def test_disabled_is_a_403_pointing_at_the_proxy_admin(self): with pytest.raises(HTTPException) as exc: - raise_for_team_admin_edit_verdict(TeamAdminEditingDisabled()) + team_admin_request_or_raise(TeamAdminEditingDisabled()) assert exc.value.status_code == 403 assert "cannot edit team settings" in exc.value.detail assert "Settings > UI > Team admin editable fields" in exc.value.detail def test_field_not_permitted_is_a_403_naming_the_field(self): with pytest.raises(HTTPException) as exc: - raise_for_team_admin_edit_verdict(TeamAdminFieldNotPermitted(field="blocked")) + team_admin_request_or_raise(TeamAdminFieldNotPermitted(field="blocked")) assert exc.value.status_code == 403 assert "'blocked'" in exc.value.detail 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 f507311a24f..3dfd994bcee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15078,6 +15078,42 @@ async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enabl assert "'rpm_limit'" in str(refused.value.message) +@pytest.mark.asyncio +async def test_update_team_team_admin_resending_budget_settings_does_not_push_back_budget_resets( + disable_audit_logging_for_mocked_team, +): + """A resent budget_duration or budget_limits would otherwise recompute the reset timestamps from now.""" + import contextlib + + stored_windows = [{"budget_duration": "7d", "max_budget": 5.0, "reset_at": "2026-09-20T00:00:00Z"}] + budgeted_team = MagicMock() + budgeted_team.metadata = {} + budgeted_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "metadata": {}, + "budget_duration": "30d", + "budget_limits": stored_windows, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=budgeted_team) + stack.enter_context(_team_admin_may_edit("tpm_limit")) + await update_team( + data=UpdateTeamRequest( + team_id="test_team_id", tpm_limit=5000, budget_duration="30d", budget_limits=stored_windows + ), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + written = prisma.db.litellm_teamtable.update.call_args.kwargs["data"] + assert written["tpm_limit"] == 5000 + assert not {"budget_duration", "budget_reset_at", "budget_limits"} & written.keys() + + @pytest.mark.asyncio async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit_logging_for_mocked_team): """The org ceiling lives on the org's budget row, so /team/update must load it to enforce the cap.""" From 02bccfd89f0bdfdb81c4c28862cc6fc61ed55443 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:13:05 +0000 Subject: [PATCH 78/89] fix(streaming): fill text_tokens when reasoning is counted from stream content Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 3 ++- .../test_streaming_chunk_builder_cursor.py | 26 ++++--------------- 2 files changed, 7 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index a8b1f81702c..99c02cdfc6a 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1100,7 +1100,8 @@ class ChunkProcessor: if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens + reasoning_tokens=reasoning_tokens, + text_tokens=max(0, returned_usage.completion_tokens - reasoning_tokens), ) elif ( returned_usage.completion_tokens_details is not None diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index f4dbb28533f..e925fd9b4a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -257,14 +257,6 @@ class TestAnthropicCursorBug: @pytest.mark.parametrize("placeholder", [1, 3, 8]) def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int): - """ - message_start placeholders are not always 1 (live Anthropic streams - have been observed sending 1 and 8 for the same prompt), and a thinking - model cut off before message_delta has streamed only reasoning_content. - The recovered usage, including the completion_tokens_details the cost - calculator bills from, must come from that reasoning rather than from - the placeholder. - """ message_start = _make_chunk( usage=Usage( prompt_tokens=100, @@ -292,13 +284,9 @@ class TestAnthropicCursorBug: ) assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens details = response.usage.completion_tokens_details - assert (details.text_tokens or 0) + details.reasoning_tokens == response.usage.completion_tokens + assert details.text_tokens + details.reasoning_tokens == response.usage.completion_tokens def test_fallback_counts_reasoning_and_text_together(self): - """ - With no usable provider count, the estimate covers everything the - provider generated: reasoning_content plus visible text, not text alone. - """ reasoning = "First I should check whether the input is sorted. " * 10 text = "The list is already sorted, so no work is needed." chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)] @@ -306,16 +294,12 @@ class TestAnthropicCursorBug: response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}]) text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True) - reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens - assert reasoning_tokens > 0 - assert response.usage.completion_tokens == text_only + reasoning_tokens + details = response.usage.completion_tokens_details + assert details.reasoning_tokens > 0 + assert response.usage.completion_tokens == text_only + details.reasoning_tokens + assert details.text_tokens == text_only def test_lone_usage_event_with_finish_reason_is_trusted(self): - """ - Guardrails rebuild responses from the chunks yielded to the client, - which excludes the un-yielded message_start. A finished stream then has - exactly one usage event (message_delta) and it must be kept as-is. - """ chunks = [ _make_chunk(content="Yes, "), _make_chunk(content="that works."), From a1ad95dbbd5dc9003598278c610b401f797d6403 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:13:41 -0700 Subject: [PATCH 79/89] fix(gemini): read the minimal thinking floor from the cost map and cover the /v1/messages bridge --- .../llms/openai/chat/gpt_5_transformation.py | 4 +- .../vertex_and_google_ai_studio_gemini.py | 6 ++- ...odel_prices_and_context_window_backup.json | 6 +++ litellm/utils.py | 6 +-- model_prices_and_context_window.json | 6 +++ .../llms/openai/test_gpt5_transformation.py | 16 +++---- ...test_vertex_and_google_ai_studio_gemini.py | 47 +++++++++++++++++-- 7 files changed, 71 insertions(+), 20 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index b02f953425d..1b93df95341 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -4,9 +4,9 @@ from typing import Final import litellm from litellm.utils import ( - _is_explicitly_disabled_factory, _supports_factory, declared_value_factory, + is_explicitly_disabled_factory, ) from .gpt_transformation import OpenAIGPTConfig @@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): Use this for opt-out checks where unknown models should be allowed through. """ - return _is_explicitly_disabled_factory( + return is_explicitly_disabled_factory( model=cls._model_map_lookup_name(model), custom_llm_provider=None, key=f"supports_{level}_reasoning_effort", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d719d53e19f..7d616c37ec1 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -79,6 +79,7 @@ from litellm.utils import ( CustomStreamWrapper, ModelResponse, is_base64_encoded, + is_explicitly_disabled_factory, supports_reasoning, ) @@ -110,7 +111,6 @@ else: SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") -GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING: Final = ("gemini-3.7-flash", "gemini-3.8-flash") def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: @@ -865,7 +865,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _supports_minimal_thinking_level(model: str) -> bool: lowered: Final = model.lower() is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered - return is_gemini3flash and not any(m in lowered for m in GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING) + return is_gemini3flash and not is_explicitly_disabled_factory( + model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort" + ) @staticmethod def _map_reasoning_effort_to_thinking_level( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f91cf82f41..43399d5af53 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25159,6 +25159,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -25216,6 +25217,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27081,6 +27083,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27140,6 +27143,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27555,6 +27559,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27612,6 +27617,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..b6f4e85a702 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2675,7 +2675,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str """Return a string value the model map declares for *key*, or ``None`` when it says nothing. The string-valued sibling of :func:`_supports_factory` and - :func:`_is_explicitly_disabled_factory`, public where those two are not because it is read + :func:`is_explicitly_disabled_factory`, public like the latter because both are read from the provider configs rather than from this module, sharing their ``get_llm_provider`` -> ``_get_model_info_helper`` chain and their unprefixed-twin fallback (#20885), so a provider-prefixed entry that omits the key still answers @@ -2711,7 +2711,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str return None -def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool: +def is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool: """Return True only when the model map explicitly sets *key* to ``False``. This is the opt-out mirror of :func:`_supports_factory`. Where @@ -2830,7 +2830,7 @@ def is_vision_explicitly_disabled(model: str, custom_llm_provider: str | None = The opt-out mirror of :func:`supports_vision`: a missing declaration reads as not disabled, so unknown or newly added models stay eligible for image routing. """ - return _is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision") + return is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision") def supports_vision(model: str, custom_llm_provider: str | None = None) -> bool: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f91cf82f41..43399d5af53 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25159,6 +25159,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -25216,6 +25217,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27081,6 +27083,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27140,6 +27143,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27555,6 +27559,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27612,6 +27617,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index b538fad71a2..ba51209e0d5 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.openai import OpenAIConfig from litellm.utils import ( - _is_explicitly_disabled_factory, + is_explicitly_disabled_factory, peek_reasoning_summary_aliases, strip_reasoning_summary_aliases_from_optional_params, ) @@ -524,19 +524,19 @@ def test_gpt5_minimal_explicitly_disabled_check(gpt5_config: OpenAIGPT5Config): def test_is_explicitly_disabled_factory_minimal(): - """_is_explicitly_disabled_factory returns True only for explicit False entries. + """is_explicitly_disabled_factory returns True only for explicit False entries. Verifies the shared helper used by _is_reasoning_effort_level_explicitly_disabled directly — so future changes to the helper are caught without going through the method wrapper. """ key = "supports_minimal_reasoning_effort" - assert _is_explicitly_disabled_factory("gpt-5.4-mini", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4-nano", None, key) - assert _is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4-pro", None, key) - assert not _is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-mini", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-nano", None, key) + assert is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key) + assert is_explicitly_disabled_factory("gpt-5.4", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-pro", None, key) + assert not is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key) def test_gpt5_unknown_model_passes_through_minimal(config: OpenAIConfig): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index a1c31689d09..b36c4e6c205 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5,11 +5,14 @@ from copy import deepcopy from typing import Final, List, cast from unittest.mock import MagicMock, patch +import httpx import pytest from pydantic import BaseModel import litellm from litellm import ModelResponse, completion +from litellm.llms.anthropic.experimental_pass_through.messages import handler as anthropic_messages_handler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -2683,7 +2686,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): [ "gemini-3.7-flash", "vertex_ai/gemini-3.8-flash", - "gemini-3.8-flash-preview", + "gemini/gemini-3.8-flash", ], ) @pytest.mark.parametrize( @@ -2691,7 +2694,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): [("minimal", True), ("none", False), ("disable", False)], ) def test_gemini_37_38_flash_floor_minimal_thinking_level( - model, reasoning_effort, include_thoughts + local_model_cost_map, model, reasoning_effort, include_thoughts ): result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( reasoning_effort, model @@ -2717,7 +2720,7 @@ def test_gemini_37_38_flash_floor_minimal_thinking_level( ], ) def test_gemini_flash_minimal_thinking_support( - model, reasoning_effort, expected_level, include_thoughts + local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts ): result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( reasoning_effort, model @@ -2727,7 +2730,7 @@ def test_gemini_flash_minimal_thinking_support( assert result["includeThoughts"] is include_thoughts -def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): +def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) thinking_param = {"type": "enabled", "budget_tokens": 1024} @@ -2742,7 +2745,7 @@ def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): assert result_36["thinkingLevel"] == "minimal" -def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): +def test_gemini_38_flash_public_reasoning_effort_none_uses_low(local_model_cost_map): result = VertexGeminiConfig().map_openai_params( non_default_params={"reasoning_effort": "none"}, optional_params={}, @@ -2756,6 +2759,40 @@ def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): } +@pytest.mark.asyncio +async def test_gemini_38_flash_messages_bridge_thinking_disabled_sends_low_thinking_level(local_model_cost_map): + captured: dict[str, dict] = {} + + def upstream(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + + await anthropic_messages_handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="gemini/gemini-3.8-flash", + custom_llm_provider="gemini", + thinking={"type": "disabled"}, + api_key="fake-gemini-key", + client=client, + ) + + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "low", + "includeThoughts": False, + } + + def test_reasoning_effort_dict_format_gemini_3(): """ Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. From 8ae1f763394bcd5a74cf3bf76ccd3756399d1958 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 22:36:56 +0000 Subject: [PATCH 80/89] feat(rust): scaffold redis cache crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 55 ++++++++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-redis/Cargo.toml | 11 ++ litellm-rust/crates/cache-redis/src/cache.rs | 127 ++++++++++++++++++ litellm-rust/crates/cache-redis/src/lib.rs | 3 + .../crates/cache-redis/tests/cache.rs | 6 + 6 files changed, 203 insertions(+) create mode 100644 litellm-rust/crates/cache-redis/Cargo.toml create mode 100644 litellm-rust/crates/cache-redis/src/cache.rs create mode 100644 litellm-rust/crates/cache-redis/src/lib.rs create mode 100644 litellm-rust/crates/cache-redis/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7397742369b..d5b94e261e6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -70,6 +70,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + [[package]] name = "async-compression" version = "0.4.46" @@ -1915,6 +1921,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "redis", + "serde_json", +] + [[package]] name = "litellm-core" version = "0.1.0" @@ -2140,6 +2155,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -2656,6 +2681,24 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redis" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" +dependencies = [ + "arcstr", + "combine", + "itoa", + "num-bigint", + "percent-encoding", + "ryu", + "sha1_smol", + "socket2 0.6.5", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3096,6 +3139,12 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -4182,6 +4231,12 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "xxhash-rust" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" + [[package]] name = "yoke" version = "0.8.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 879090870d8..eb3413eb89a 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -23,6 +23,7 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" +redis = "1.7.0" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml new file mode 100644 index 00000000000..2db2ae33840 --- /dev/null +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-cache-redis" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +redis.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs new file mode 100644 index 00000000000..963783d31a6 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -0,0 +1,127 @@ +use std::sync::Mutex; +use std::time::Duration; + +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, + Error, +}; +use redis::Commands; + +const DEFAULT_TTL: Duration = Duration::from_secs(600); + +pub struct RedisCache { + connection: Mutex, + default_ttl: Duration, +} + +impl RedisCache { + pub fn new(url: &str, default_ttl: Option) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let connection = client.get_connection().map_err(|_| Error::Unavailable)?; + Ok(Self { + connection: Mutex::new(connection), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + }) + } + + fn connection(&self) -> Result, Error> { + self.connection.lock().map_err(|_| Error::Unavailable) + } + + fn encode(value: &CacheEntry) -> Result, Error> { + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(value: Vec) -> Result { + serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry) + } + + fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs().max(1) + } +} + +impl BaseCache for RedisCache { + type Value = CacheEntry; + + fn default_ttl(&self) -> Duration { + self.default_ttl + } + + fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { + let payload = Self::encode(&value)?; + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + self.connection()? + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable) + } + + fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { + self.connection()? + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable)? + .map(Self::decode) + .transpose() + } + + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.connection()? + .del::<_, ()>(key) + .map_err(|_| Error::Unavailable) + } + + fn flush_cache(&self) -> Result<(), Error> { + self.connection()? + .flushdb::<()>() + .map_err(|_| Error::Unavailable) + } + + fn disconnect(&self) -> CacheFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } + + fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { + Box::pin(async { + let mut connection = self.connection()?; + redis::cmd("PING") + .query::(&mut *connection) + .map_err(|_| Error::Unavailable)?; + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }) + }) + } +} + +#[cfg(test)] +mod tests { + use super::RedisCache; + use litellm_cache::CacheEntry; + use serde_json::json; + use std::time::Duration; + + #[test] + fn cache_entries_round_trip_through_json() { + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + + let encoded = RedisCache::encode(&entry).unwrap(); + assert_eq!(RedisCache::decode(encoded).unwrap(), entry); + } + + #[test] + fn invalid_json_is_rejected() { + assert!(RedisCache::decode(b"not json".to_vec()).is_err()); + } + + #[test] + fn ttl_seconds_keeps_redis_expiration_positive() { + assert_eq!(RedisCache::ttl_seconds(Duration::ZERO), 1); + assert_eq!(RedisCache::ttl_seconds(Duration::from_millis(1500)), 1); + assert_eq!(RedisCache::ttl_seconds(Duration::from_secs(15)), 15); + } +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs new file mode 100644 index 00000000000..37b35c5ea4a --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -0,0 +1,3 @@ +mod cache; + +pub use cache::RedisCache; diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs new file mode 100644 index 00000000000..76f73145da8 --- /dev/null +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -0,0 +1,6 @@ +use litellm_cache_redis::RedisCache; + +#[test] +fn constructor_rejects_invalid_urls() { + assert!(RedisCache::new("not a redis url", None).is_err()); +} From 0b3c3885bced09b4e98aaed57b58fca742a20339 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 22:38:38 +0000 Subject: [PATCH 81/89] fix(rust): scope Redis dependency to cache crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 1 - litellm-rust/crates/cache-redis/Cargo.toml | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index eb3413eb89a..879090870d8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -23,7 +23,6 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" -redis = "1.7.0" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 2db2ae33840..d954084168f 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,5 +7,5 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis.workspace = true +redis = "1.7.0" serde_json.workspace = true From 2a9fa48730bf4cfe19b5a36d5144d11944b42a5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:18:44 -0700 Subject: [PATCH 82/89] fix: expand wildcard deployments for proxy admins on /model_group/info --- litellm/proxy/proxy_server.py | 10 +++++++++- .../proxy/proxy_server/test_routes_model_info.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6fe2ab57112..2c5c635b14b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15542,7 +15542,15 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - llm_router.get_model_names() + get_complete_model_list( + key_models=(), + team_models=(), + proxy_model_list=llm_router.get_model_names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) if is_proxy_admin else await get_available_models_for_user( user_api_key_dict=user_api_key_dict, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 2c101156c6c..3000a2ea101 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -353,6 +353,22 @@ def test_model_group_info_proxy_admin_ignores_key_model_restriction( assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"] +@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) +def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model_group_info_router.get_model_names.return_value = ["gpt-4", "anthropic/*"] + known_anthropic_models = get_known_models_from_wildcard(wildcard_model="anthropic/*") + assert known_anthropic_models + + with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", *known_anthropic_models] + + def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router): from litellm.proxy._types import LitellmUserRoles From 93ba409adf9215ca051372714add07a5c406f89e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 23:19:28 +0000 Subject: [PATCH 83/89] fix(rust): address Redis cache review findings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 46 ++++ litellm-rust/crates/cache-redis/Cargo.toml | 4 + litellm-rust/crates/cache-redis/src/cache.rs | 252 ++++++++++++++++--- 3 files changed, 270 insertions(+), 32 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d5b94e261e6..9cfce7e0704 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1843,6 +1843,12 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litellm-auth" version = "0.1.0" @@ -1927,7 +1933,9 @@ version = "0.1.0" dependencies = [ "litellm-cache", "redis", + "redis-test", "serde_json", + "tokio", ] [[package]] @@ -2699,6 +2707,18 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "redis-test" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca" +dependencies = [ + "rand 0.9.5", + "redis", + "socket2 0.6.5", + "tempfile", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2889,6 +2909,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.21.12" @@ -3348,6 +3381,19 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "1.0.69" diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index d954084168f..933b0feaae4 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -9,3 +9,7 @@ repository.workspace = true litellm-cache.workspace = true redis = "1.7.0" serde_json.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 963783d31a6..69dee6c6363 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,4 +1,4 @@ -use std::sync::Mutex; +use std::sync::{Arc, Mutex, MutexGuard}; use std::time::Duration; use litellm_cache::{ @@ -8,26 +8,45 @@ use litellm_cache::{ use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); +const KEY_PREFIX: &str = "litellm-cache:"; -pub struct RedisCache { - connection: Mutex, +pub struct RedisCache { + connection: Arc>, default_ttl: Duration, } -impl RedisCache { +impl RedisCache { pub fn new(url: &str, default_ttl: Option) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self { - connection: Mutex::new(connection), + Ok(Self::with_connection(connection, default_ttl)) + } +} + +impl RedisCache +where + C: redis::ConnectionLike + Send + 'static, +{ + fn with_connection(connection: C, default_ttl: Option) -> Self { + Self { + connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - }) + } } - fn connection(&self) -> Result, Error> { + fn connection(&self) -> Result, Error> { self.connection.lock().map_err(|_| Error::Unavailable) } + fn namespaced_key(key: &str) -> String { + format!("{KEY_PREFIX}{key}") + } + + fn namespaced_pattern() -> &'static str { + const PATTERN: &str = "litellm-cache:*"; + PATTERN + } + fn encode(value: &CacheEntry) -> Result, Error> { serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) } @@ -37,11 +56,31 @@ impl RedisCache { } fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs().max(1) + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) + } + + fn run_blocking(connection: Arc>, operation: F) -> CacheFuture<'static, T> + where + T: Send + 'static, + F: FnOnce(&mut C) -> Result + Send + 'static, + { + Box::pin(async move { + tokio::task::spawn_blocking(move || { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut connection) + }) + .await + .map_err(|_| Error::Unavailable)? + }) } } -impl BaseCache for RedisCache { +impl BaseCache for RedisCache +where + C: redis::ConnectionLike + Send + 'static, +{ type Value = CacheEntry; fn default_ttl(&self) -> Duration { @@ -52,13 +91,13 @@ impl BaseCache for RedisCache { let payload = Self::encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); self.connection()? - .set_ex::<_, _, ()>(key, payload, ttl) + .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl) .map_err(|_| Error::Unavailable) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { self.connection()? - .get::<_, Option>>(key) + .get::<_, Option>>(Self::namespaced_key(key)) .map_err(|_| Error::Unavailable)? .map(Self::decode) .transpose() @@ -66,26 +105,101 @@ impl BaseCache for RedisCache { fn delete_cache(&self, key: &str) -> Result<(), Error> { self.connection()? - .del::<_, ()>(key) + .del::<_, ()>(Self::namespaced_key(key)) .map_err(|_| Error::Unavailable) } fn flush_cache(&self) -> Result<(), Error> { - self.connection()? - .flushdb::<()>() + let mut connection = self.connection()?; + let keys = connection + .scan_match(Self::namespaced_pattern()) + .map_err(|_| Error::Unavailable)? + .collect::>>() + .map_err(|_| Error::Unavailable)?; + if keys.is_empty() { + return Ok(()); + } + connection + .del::<_, usize>(keys) + .map(|_| ()) .map_err(|_| Error::Unavailable) } + fn async_set_cache<'a>( + &'a self, + key: &'a str, + value: Self::Value, + kwargs: CacheKwargs, + ) -> CacheFuture<'a, ()> { + let payload = Self::encode(&value); + let key = Self::namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .set_ex::<_, _, ()>(key, payload?, ttl) + .map_err(|_| Error::Unavailable) + }) + } + + fn async_get_cache<'a>( + &'a self, + key: &'a str, + _: &'a CacheKwargs, + ) -> CacheFuture<'a, Option> { + let key = Self::namespaced_key(key); + Box::pin(async move { + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable) + }) + .await? + .map(Self::decode) + .transpose() + }) + } + + fn async_set_cache_pipeline<'a>( + &'a self, + cache_list: Vec<(String, Self::Value)>, + kwargs: CacheKwargs, + ) -> CacheFuture<'a, ()> { + let entries = cache_list + .into_iter() + .map(|(key, value)| { + Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload)) + }) + .collect::, _>>(); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + for (key, payload) in entries? { + connection + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + }) + } + + fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { + let key = Self::namespaced_key(key); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) + }) + } + fn disconnect(&self) -> CacheFuture<'_, ()> { Box::pin(async { Ok(()) }) } fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async { - let mut connection = self.connection()?; - redis::cmd("PING") - .query::(&mut *connection) - .map_err(|_| Error::Unavailable)?; + Box::pin(async move { + Self::run_blocking(Arc::clone(&self.connection), |connection| { + redis::cmd("PING") + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; Ok(CacheConnectionResult { status: CacheConnectionStatus::Success, message: "Redis cache connection test successful".into(), @@ -98,30 +212,104 @@ impl BaseCache for RedisCache { #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::CacheEntry; + use litellm_cache::{BaseCache, CacheEntry, CacheKwargs}; + use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; - #[test] - fn cache_entries_round_trip_through_json() { - let entry = CacheEntry { + fn entry() -> CacheEntry { + CacheEntry { timestamp: 123.0, response: json!({"choices": [{"text": "cached"}]}), - }; + } + } - let encoded = RedisCache::encode(&entry).unwrap(); - assert_eq!(RedisCache::decode(encoded).unwrap(), entry); + #[test] + fn cache_entries_round_trip_through_json() { + let entry = entry(); + let encoded = RedisCache::::encode(&entry).unwrap(); + assert_eq!( + RedisCache::::decode(encoded).unwrap(), + entry + ); } #[test] fn invalid_json_is_rejected() { - assert!(RedisCache::decode(b"not json".to_vec()).is_err()); + assert!(RedisCache::::decode(b"not json".to_vec()).is_err()); } #[test] - fn ttl_seconds_keeps_redis_expiration_positive() { - assert_eq!(RedisCache::ttl_seconds(Duration::ZERO), 1); - assert_eq!(RedisCache::ttl_seconds(Duration::from_millis(1500)), 1); - assert_eq!(RedisCache::ttl_seconds(Duration::from_secs(15)), 15); + fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { + assert_eq!( + RedisCache::::ttl_seconds(Duration::ZERO), + 1 + ); + assert_eq!( + RedisCache::::ttl_seconds(Duration::from_millis(1500)), + 2 + ); + assert_eq!( + RedisCache::::ttl_seconds(Duration::from_secs(15)), + 15 + ); + } + + #[test] + fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { + let value = entry(); + let payload = RedisCache::::encode(&value).unwrap(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:key") + .arg(600) + .arg(payload.clone()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + cache + .set_cache("key", value.clone(), CacheKwargs::default()) + .unwrap(); + assert_eq!( + cache.get_cache("key", &CacheKwargs::default()).unwrap(), + Some(value) + ); + cache.delete_cache("key").unwrap(); + } + + #[test] + fn flush_scans_and_deletes_only_cache_keys() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("litellm-cache:*"), + Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + cache.flush_cache().unwrap(); + } + + #[tokio::test] + async fn test_connection_runs_ping_off_executor() { + let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + assert_eq!( + cache.test_connection().await.unwrap().status, + litellm_cache::CacheConnectionStatus::Success + ); } } From ab92a6637d1308af6215edda3a50582548d639bf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:26:52 -0700 Subject: [PATCH 84/89] test(e2e): cover team admin editable fields on /team/update Team admins are refused until a proxy admin enables a field, then limited to the enabled fields, and resending unchanged budget settings keeps the team's budget reset times --- tests/e2e/coverage_registry/mgmt.yaml | 3 + .../management/test_team_management_e2e.py | 299 +++++++++++++++++- 2 files changed, 292 insertions(+), 10 deletions(-) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 31ad61ba3e2..d93d2b2cc67 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -30,6 +30,9 @@ - {id: mgmt.key.health.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4292", rationale: "Key health endpoint"} - {id: mgmt.key.bulk_update.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:2677", rationale: "Batch key updates"} - {id: mgmt.team.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1582", rationale: "Metadata/budget updates persist"} +- {id: mgmt.team.update.team_admin_forbidden_until_enabled, module: mgmt, tier: P0, surface: api, assertions: [team_admin_forbidden_until_enabled], source: "team_admin_field_permissions.py:156", rationale: "With no team admin editable fields enabled, a team admin's /team/update is 403 and /team/info reports editing disabled"} +- {id: mgmt.team.update.team_admin_limited_to_enabled_fields, module: mgmt, tier: P0, surface: api, assertions: [team_admin_limited_to_enabled_fields], source: "team_admin_field_permissions.py:156", rationale: "A team admin may change only the enabled fields; a request that also changes any other field is 403 and writes nothing"} +- {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"} - {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"} - {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"} - {id: mgmt.team.info.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "team_endpoints.py:2244", rationale: "Metadata+members+budgets"} diff --git a/tests/e2e/management/test_team_management_e2e.py b/tests/e2e/management/test_team_management_e2e.py index 108aeaad21b..f21931b6ff1 100644 --- a/tests/e2e/management/test_team_management_e2e.py +++ b/tests/e2e/management/test_team_management_e2e.py @@ -1,5 +1,6 @@ """Live e2e: the /team/* management routes' block, membership, and admin-only -contract. +contract, plus the team settings a team admin may change on /team/update once a +proxy admin enables them under Settings > UI > Team admin editable fields. Each test creates its team/user/key resources under unique names (deleted on teardown) and asserts both halves of the contract: the recorded state (the info @@ -8,21 +9,25 @@ Team writes reach the read path once their db/cache entry propagates, so the read-backs poll to a deadline instead of asserting once. Everything the shared harness does not already model lives here: the local -request/response models for /team/block, /team/member_update, and the -/team/info fields (blocked flag and per-member budget) these tests assert on. +request/response models for /team/block, /team/member_update, the partial +/team/update, the UI settings allow-list, and the /team/info fields (blocked +flag, limits, budgets, per-member budget, the caller's edit access) these tests +assert on. """ from __future__ import annotations import time -from collections.abc import Callable -from typing import Literal +from collections.abc import Callable, Generator +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from typing import Final, Literal import pytest from pydantic import BaseModel -from e2e_config import unique_marker -from e2e_http import NoBody, StreamingResponse, unwrap +from e2e_config import settle_propagation, unique_marker +from e2e_http import NoBody, PartialBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import ( @@ -39,6 +44,8 @@ pytestmark = pytest.mark.e2e TeamRole = Literal["admin", "user"] +_TEAM_TPM_LIMIT: Final = 1000 + class TeamBlockBody(BaseModel): team_id: str @@ -66,11 +73,37 @@ class TeamMembership(BaseModel): litellm_budget_table: MemberBudgetTable | None = None -class TeamInfoData(BaseModel): +class CallerEditAccess(BaseModel): + kind: Literal["unrestricted", "team_admin", "team_admin_disabled", "none"] + editable_fields: list[str] = [] + + +class BudgetWindow(BaseModel): + budget_duration: str + max_budget: float + reset_at: str | None = None + + +class TeamCustomMetadata(BaseModel): + cost_center: str | None = None + + +class TeamSettings(BaseModel): team_alias: str | None = None models: list[str] = [] + tpm_limit: int | None = None + rpm_limit: int | None = None + max_budget: float | None = None + budget_duration: str | None = None + budget_limits: list[BudgetWindow] | None = None + metadata: TeamCustomMetadata | None = None + + +class TeamInfoData(TeamSettings): blocked: bool | None = None members_with_roles: list[MemberRoleEntry] = [] + budget_reset_at: datetime | None = None + caller_edit_access: CallerEditAccess | None = None class TeamInfoRead(BaseModel): @@ -79,6 +112,27 @@ class TeamInfoRead(BaseModel): team_memberships: list[TeamMembership] = [] +class TeamWithAdminNewBody(TeamNewBody): + tpm_limit: int + members_with_roles: list[TeamMemberEntry] + + +class TeamSettingsChange(PartialBody, TeamSettings): + pass + + +class TeamSettingsUpdate(TeamSettingsChange): + team_id: str + + +class TeamAdminEditableFields(BaseModel): + team_admin_editable_team_fields: list[str] = [] + + +class UiSettingsRead(BaseModel): + values: TeamAdminEditableFields + + def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: @@ -107,17 +161,27 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke return key -def _read_team(client: ManagementClient, team_id: str) -> TeamInfoRead: +def _read_team(client: ManagementClient, team_id: str, caller_key: str | None = None) -> TeamInfoRead: return unwrap( client.proxy.transport.get( "/team/info", - headers=client.proxy.transport.master, + headers=client.proxy.transport.master if caller_key is None else client.proxy.transport.bearer(caller_key), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoRead, ) ) +def _poll_team( + client: ManagementClient, team_id: str, ready: Callable[[TeamInfoData], bool], failure: str +) -> TeamInfoData: + def read() -> TeamInfoData | None: + info = _read_team(client, team_id).team_info + return info if ready(info) else None + + return _poll(client, read, failure) + + def _set_blocked(client: ManagementClient, team_id: str, *, blocked: bool) -> None: _ = unwrap( client.proxy.transport.post( @@ -301,3 +365,218 @@ class TestTeamManagementRoutes: client.add_team_member(team_id, member_id) member_key = _generate_key(client, resources, KeyGenerateBody(user_id=member_id, team_id=team_id)) return member_id, other_id, member_key, team_id + + +def _team_admin_editable_fields(client: ManagementClient) -> list[str]: + return unwrap( + client.proxy.transport.get( + "/get/ui_settings", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=UiSettingsRead, + ) + ).values.team_admin_editable_team_fields + + +def _set_team_admin_editable_fields(client: ManagementClient, fields: list[str]) -> None: + _ = unwrap( + client.proxy.transport.patch( + "/update/ui_settings", + headers=client.proxy.transport.master, + json=TeamAdminEditableFields(team_admin_editable_team_fields=fields), + response_type=NoBody, + ) + ) + + +@contextmanager +def _team_admins_may_edit(client: ManagementClient, fields: list[str]) -> Generator[None]: + """The allow-list is proxy-wide, so restore whatever was there. Other replicas pick a change up on their + config reload, which the wait covers before any team admin call lands on one of them.""" + original = _team_admin_editable_fields(client) + _set_team_admin_editable_fields(client, fields) + settle_propagation(time.monotonic()) + try: + yield + finally: + _set_team_admin_editable_fields(client, original) + + +@pytest.fixture(scope="class") +def no_team_admin_editable_fields(client: ManagementClient) -> Generator[None]: + with _team_admins_may_edit(client, []): + yield + + +@pytest.fixture(scope="class") +def tpm_limit_editable_by_team_admins(client: ManagementClient) -> Generator[None]: + with _team_admins_may_edit(client, ["tpm_limit"]): + yield + + +def _team_with_admin(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]: + """A team with a tpm_limit, and the key of a user who is an admin of that team.""" + admin_id = _create_user(client, resources, f"e2e-team-admin-{unique_marker()}@example.com") + team_id = client.create_team( + TeamWithAdminNewBody( + team_alias=f"e2e-team-admin-{unique_marker()}", + tpm_limit=_TEAM_TPM_LIMIT, + members_with_roles=[TeamMemberEntry(role="admin", user_id=admin_id)], + ) + ) + resources.defer(lambda: client.delete_team(team_id)) + return team_id, _generate_key(client, resources, KeyGenerateBody(user_id=admin_id)) + + +def _update_team_as(client: ManagementClient, caller_key: str, body: TeamSettingsUpdate) -> StreamingResponse: + return client.proxy.transport.send("/team/update", headers=client.proxy.transport.bearer(caller_key), json=body) + + +@pytest.mark.usefixtures("no_team_admin_editable_fields") +class TestTeamAdminWithNoEditableFields: + """No proxy admin has enabled a team field for team admins, which is how every proxy starts.""" + + @pytest.mark.covers("mgmt.team.update.team_admin_forbidden_until_enabled") + def test_team_admin_cannot_change_any_team_setting( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + team_id, admin_key = _team_with_admin(client, resources) + access = _read_team(client, team_id, admin_key).team_info.caller_edit_access + assert access == CallerEditAccess(kind="team_admin_disabled"), ( + f"/team/info should tell the team admin that editing is disabled, got {access}" + ) + + outcome = _update_team_as(client, admin_key, TeamSettingsUpdate(team_id=team_id, tpm_limit=5000)) + + assert outcome.status_code == 403, ( + f"/team/update by a team admin must be 403 while nothing is enabled, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + assert "cannot edit team settings" in outcome.body, f"403 body should say why, got: {outcome.body[:300]}" + tpm_limit = _read_team(client, team_id).team_info.tpm_limit + assert tpm_limit == _TEAM_TPM_LIMIT, f"the refused update still changed tpm_limit to {tpm_limit}" + + +@pytest.mark.usefixtures("tpm_limit_editable_by_team_admins") +class TestTeamAdminWithTpmLimitEnabled: + """A proxy admin has enabled tpm_limit, so a team admin may change that setting and no other.""" + + @pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields") + def test_team_admin_saves_the_settings_form_with_a_new_tpm_limit( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + team_id, admin_key = _team_with_admin(client, resources) + access = _read_team(client, team_id, admin_key).team_info.caller_edit_access + assert access == CallerEditAccess(kind="team_admin", editable_fields=["tpm_limit"]), ( + f"/team/info should list tpm_limit as the team admin's only editable field, got {access}" + ) + before = _read_team(client, team_id).team_info + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate(team_id=team_id, team_alias=before.team_alias, models=before.models, tpm_limit=5000), + ) + + assert outcome.status_code == 200, ( + f"a team admin resending the form with only tpm_limit changed must succeed, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + after = _poll_team( + client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000" + ) + assert after.model_copy(update={"tpm_limit": _TEAM_TPM_LIMIT}) == before, ( + f"the update changed more than tpm_limit: before {before}, after {after}" + ) + + @pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields") + @pytest.mark.parametrize( + "change", + [ + pytest.param(TeamSettingsChange(rpm_limit=10), id="rpm_limit"), + pytest.param(TeamSettingsChange(max_budget=0.5), id="max_budget"), + pytest.param(TeamSettingsChange(team_alias="renamed-by-team-admin"), id="team_alias"), + pytest.param(TeamSettingsChange(models=["gemini-2.5-flash"]), id="models"), + pytest.param(TeamSettingsChange(budget_duration="1d"), id="budget_duration"), + pytest.param(TeamSettingsChange(metadata=TeamCustomMetadata(cost_center="team-admin")), id="metadata"), + ], + ) + def test_team_admin_cannot_change_a_setting_that_is_not_enabled( + self, client: ManagementClient, resources: ResourceManager, change: TeamSettingsChange + ) -> None: + (field,) = change.model_fields_set + team_id, admin_key = _team_with_admin(client, resources) + before = _read_team(client, team_id).team_info + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate.model_validate( + {**change.model_dump(exclude_unset=True), "team_id": team_id, "tpm_limit": 5000} + ), + ) + + assert outcome.status_code == 403, ( + f"a team admin changing {field} must be 403, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert f"'{field}'" in outcome.body, f"403 body should name {field}, got: {outcome.body[:300]}" + after = _read_team(client, team_id).team_info + assert after == before, ( + f"the refused update still wrote to the team, the enabled tpm_limit included: before {before}, " + f"after {after}" + ) + + @pytest.mark.covers("mgmt.team.update.team_admin_resend_keeps_budget_reset") + def test_team_admin_resending_the_budget_settings_keeps_the_next_budget_reset( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """A 120s budget resets at the start of the minute after next. Resending it once the next minute has + started would push that reset a minute later, while the stored reset is still a minute out, so the + proxy's budget reset job cannot be what moves it.""" + team_id, admin_key = _team_with_admin(client, resources) + _ = unwrap( + client.proxy.transport.post( + "/team/update", + headers=client.proxy.transport.master, + json=TeamSettingsUpdate( + team_id=team_id, + budget_duration="120s", + budget_limits=[BudgetWindow(budget_duration="120s", max_budget=5.0)], + ), + response_type=NoBody, + ) + ) + budgeted = _poll_team( + client, + team_id, + lambda info: info.budget_reset_at is not None and bool(info.budget_limits), + "/team/info never reflected the 120s budget the proxy admin set", + ) + assert budgeted.budget_reset_at is not None + next_minute = budgeted.budget_reset_at - timedelta(seconds=58) + time.sleep(max(0.0, (next_minute - datetime.now(UTC)).total_seconds())) + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate( + team_id=team_id, + tpm_limit=5000, + budget_duration=budgeted.budget_duration, + budget_limits=budgeted.budget_limits, + ), + ) + + assert outcome.status_code == 200, ( + f"resending unchanged budget settings with a new tpm_limit must succeed, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + after = _poll_team( + client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000" + ) + assert after.budget_reset_at == budgeted.budget_reset_at, ( + f"the team admin pushed the budget reset from {budgeted.budget_reset_at} to {after.budget_reset_at}" + ) + assert after.budget_limits == budgeted.budget_limits, ( + f"the team admin pushed the budget window resets from {budgeted.budget_limits} to {after.budget_limits}" + ) From 884e96958c31c792a373b3e973b8cd20c338d3a9 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:27:48 +0000 Subject: [PATCH 85/89] fix(streaming): cap estimated reasoning tokens to the provider total and cover dict chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 20 ++--- .../test_streaming_chunk_builder_cursor.py | 81 +++++++++++++------ 2 files changed, 68 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 99c02cdfc6a..9bb2d489f83 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -5,7 +5,6 @@ from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast -from pydantic import BaseModel from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger @@ -974,12 +973,10 @@ class ChunkProcessor: return None @staticmethod - def _chunk_choices(chunk: "_UsageBearingChunk | BaseModel") -> Sequence[object]: + def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]: if isinstance(chunk, dict): return chunk.get("choices", ()) - if isinstance(chunk, (ModelResponse, ModelResponseStream)): - return chunk.choices - return () + return getattr(chunk, "choices", ()) @staticmethod def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool: @@ -1099,19 +1096,22 @@ class ChunkProcessor: if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: + capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens, - text_tokens=max(0, returned_usage.completion_tokens - reasoning_tokens), + reasoning_tokens=capped_reasoning_tokens, + text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens, ) elif ( returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) - returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens + existing_capped_reasoning_tokens: Final = min( + max(0, reasoning_tokens), returned_usage.completion_tokens + ) + returned_usage.completion_tokens_details.reasoning_tokens = existing_capped_reasoning_tokens if returned_usage.completion_tokens_details.text_tokens is None: returned_usage.completion_tokens_details.text_tokens = ( - returned_usage.completion_tokens - capped_reasoning_tokens + returned_usage.completion_tokens - existing_capped_reasoning_tokens ) if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index e925fd9b4a8..8617c5b81e8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -19,7 +19,6 @@ to 0 when the only update we saw was the cursor, allowing the text-based fallback to estimate from the real completion text. """ - import pytest import litellm @@ -71,9 +70,7 @@ class TestAnthropicCursorBug: token_counter fallback can estimate from completion text. """ # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) # Several content_block_delta chunks (no usage attached) text_chunks = [ _make_chunk(content="Hello"), @@ -99,9 +96,7 @@ class TestAnthropicCursorBug: Normal complete stream: message_start cursor=1, then message_delta=3847. Last-wins must give 3847 (the real value). """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]] # message_delta with the real cumulative output_tokens message_delta = _make_chunk( @@ -121,19 +116,14 @@ class TestAnthropicCursorBug: End-to-end via calculate_usage(): cursor-only stream + real completion text should produce a token-counter estimate, NOT 1. """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark) text_chunks = [ _make_chunk(content="Based on your question, I think the answer is "), _make_chunk(content="forty-two. Here is my reasoning: "), ] chunks = [message_start, *text_chunks] - completion_output = ( - "Based on your question, I think the answer is forty-two. " - "Here is my reasoning: " - ) + completion_output = "Based on your question, I think the answer is forty-two. Here is my reasoning: " processor = ChunkProcessor(chunks=chunks, messages=[]) usage = processor.calculate_usage( @@ -151,9 +141,7 @@ class TestAnthropicCursorBug: def test_cache_fields_preserved_from_message_start(self): """cache_read / cache_creation come from message_start and must survive.""" - message_start_usage = Usage( - prompt_tokens=1024, completion_tokens=1, total_tokens=1025 - ) + message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) # Anthropic puts these in message_start message_start_usage.cache_read_input_tokens = 512 message_start_usage.cache_creation_input_tokens = 128 @@ -195,9 +183,7 @@ class TestAnthropicCursorBug: on a 1-token string also gives ~1, so billing is still approximately correct. This test pins that the result is sane (1 or 0). """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21)) text_chunk = _make_chunk(content="Yes.") # Anthropic's message_delta also gives output_tokens=1 in this case message_delta = _make_chunk( @@ -233,9 +219,7 @@ class TestAnthropicCursorBug: must fire so token_counter estimates from completion text instead of billing the placeholder. """ - message_start_usage = Usage( - prompt_tokens=1024, completion_tokens=1, total_tokens=1025 - ) + message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) message_start_usage.cache_read_input_tokens = 4096 message_start = _make_chunk(usage=message_start_usage) # Subsequent chunks with cache fields but no completion_tokens @@ -312,6 +296,57 @@ class TestAnthropicCursorBug: result = processor._calculate_usage_per_chunk(chunks=chunks) assert result["completion_tokens"] == 5 + def test_dict_chunks_with_finish_reason_are_trusted(self): + chunks = [ + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "Yes, "}, "finish_reason": None}], + }, + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "that works."}, "finish_reason": "stop"}], + "usage": Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + }, + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + + def test_dict_chunks_without_finish_reason_reset_placeholder(self): + chunks = [ + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [], + "usage": Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21), + }, + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "partial"}, "finish_reason": None}], + }, + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 0 + assert result["completion_tokens_details"] is None + + def test_estimated_reasoning_is_capped_to_trusted_completion_total(self): + chunks = [ + _make_chunk(reasoning_content="Let me reason about this carefully and at length. " * 20), + _make_chunk( + finish_reason="stop", + usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + ), + ] + response = litellm.stream_chunk_builder( + chunks=chunks, + messages=[{"role": "user", "content": "Go."}], + ) + details = response.usage.completion_tokens_details + assert response.usage.completion_tokens == 5 + assert details.reasoning_tokens <= response.usage.completion_tokens + assert details.reasoning_tokens + details.text_tokens == response.usage.completion_tokens + assert details.text_tokens >= 0 + class TestProviderGuard: """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic From aea9678f6198002fc6f914733320d921a5aff558 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:28:28 +0000 Subject: [PATCH 86/89] refactor(streaming): compute the reasoning token cap once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_chunk_builder_utils.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 9bb2d489f83..0ca93fe08b3 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1095,8 +1095,8 @@ class ChunkProcessor: returned_usage.completion_tokens_details = completion_tokens_details if reasoning_tokens is not None: + capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) if returned_usage.completion_tokens_details is None: - capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( reasoning_tokens=capped_reasoning_tokens, text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens, @@ -1105,13 +1105,10 @@ class ChunkProcessor: returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - existing_capped_reasoning_tokens: Final = min( - max(0, reasoning_tokens), returned_usage.completion_tokens - ) - returned_usage.completion_tokens_details.reasoning_tokens = existing_capped_reasoning_tokens + returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens if returned_usage.completion_tokens_details.text_tokens is None: returned_usage.completion_tokens_details.text_tokens = ( - returned_usage.completion_tokens - existing_capped_reasoning_tokens + returned_usage.completion_tokens - capped_reasoning_tokens ) if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details From 41737aeda8d919b82c8018c2789e29211a37c4be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:37:32 -0700 Subject: [PATCH 87/89] test(proxy): drop the docstring from the agent-runtime passthrough regression class --- .../pass_through_endpoints/test_llm_pass_through_endpoints.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 7f044486e14..6e82c90514d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1984,8 +1984,6 @@ class TestBedrockAgentRuntimePassthroughToggle: class TestBedrockAgentRuntimePassthroughVirtualKeyLeak: - """Regression for LIT-7912: the agent-runtime branch of ``/bedrock/{endpoint}`` forwarded every caller header, - so a LiteLLM key presented in ``x-api-key`` or ``x-litellm-api-key`` rode to AWS next to the SigV4 signature.""" VKEY: Final = "sk-litellm-victim-key" MASTER_KEY: Final = "sk-master-1234" From 3ee8d43fdd63e068ee876a1999cf6412265fe49c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:37:35 -0700 Subject: [PATCH 88/89] fix(ui): send only a changed TPM limit from the team admin settings form Save stays disabled until the value differs from the team's, so an unchanged form never reaches /team/update --- .../team/TeamAdminSettingsForm.test.tsx | 13 ++++++++++++ .../components/team/TeamAdminSettingsForm.tsx | 7 +++++-- .../team/teamAdminEditAccess.test.ts | 20 ++++++++++++++----- .../components/team/teamAdminEditAccess.ts | 8 +++++++- 4 files changed, 40 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx index 65f7a13652d..677c8859eb2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx @@ -53,6 +53,18 @@ describe("TeamAdminSettingsForm", () => { await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: null })); }); + it("keeps Save disabled until the TPM limit differs from the team's", () => { + renderForm(new Set(["tpm_limit"])); + const tpmInput = screen.getByLabelText("Tokens per minute Limit (TPM)"); + const save = screen.getByRole("button", { name: /save changes/i }); + + expect(save).toBeDisabled(); + fireEvent.change(tpmInput, { target: { value: "5000" } }); + expect(save).toBeEnabled(); + fireEvent.change(tpmInput, { target: { value: "1000" } }); + expect(save).toBeDisabled(); + }); + it("closes without saving on cancel", async () => { const user = userEvent.setup(); const { onSave, onCancel } = renderForm(new Set(["tpm_limit"])); @@ -65,6 +77,7 @@ describe("TeamAdminSettingsForm", () => { it("locks both buttons while a save is in flight", () => { renderForm(new Set(["tpm_limit"]), { isSaving: true }); + fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } }); expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled(); expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled(); diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx index 581215b7fdf..140533fada5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx @@ -1,6 +1,7 @@ "use client"; import { Save } from "lucide-react"; +import { useWatch } from "react-hook-form"; import { z } from "zod/v4"; import { FormField } from "@/components/shared/form/FormField"; @@ -37,7 +38,9 @@ export default function TeamAdminSettingsForm({ onSave, }: TeamAdminSettingsFormProps) { const form = useZodForm(teamAdminSettingsSchema, { defaultValues: initialValues }); - const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, editableFields))); + const draft = useWatch({ control: form.control }); + const hasChanges = Object.keys(teamAdminSettingsChanges(draft, initialValues, editableFields)).length > 0; + const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, initialValues, editableFields))); return (
void submit(event)}> @@ -56,7 +59,7 @@ export default function TeamAdminSettingsForm({ - diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index c8117800053..ded6d775ce8 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -20,21 +20,31 @@ describe("teamAdminFieldLabel", () => { describe("teamAdminSettingsChanges", () => { const tpmEnabled = new Set(["tpm_limit"]); + const stored = { tpm_limit: 1000 }; it.each([ ["a typed number string", "5000", 5000], - ["a stored number", 1200, 1200], + ["a number", 1200, 1200], ["zero", "0", 0], ["an emptied input", "", null], ["whitespace", " ", null], - ["no stored limit", null, null], + ["no limit", null, null], ["an unset value", undefined, null], - ])("sends tpm_limit for %s", (_label, tpm_limit, expected) => { - expect(teamAdminSettingsChanges({ tpm_limit }, tpmEnabled)).toStrictEqual({ tpm_limit: expected }); + ])("sends tpm_limit changed to %s", (_label, tpm_limit, expected) => { + expect(teamAdminSettingsChanges({ tpm_limit }, stored, tpmEnabled)).toStrictEqual({ tpm_limit: expected }); + }); + + it.each([ + ["the stored number", 1000, { tpm_limit: 1000 }], + ["the stored number typed back in", "1000", { tpm_limit: 1000 }], + ["an emptied input over no stored limit", "", { tpm_limit: null }], + ["an unset value over no stored limit", undefined, { tpm_limit: null }], + ])("sends nothing for %s", (_label, tpm_limit, initialValues) => { + expect(teamAdminSettingsChanges({ tpm_limit }, initialValues, tpmEnabled)).toStrictEqual({}); }); it("leaves tpm_limit out when the proxy did not enable it for team admins", () => { - expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, new Set(["max_budget"]))).toStrictEqual({}); + expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, stored, new Set(["max_budget"]))).toStrictEqual({}); }); }); diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index d38566eefda..73129923907 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -59,8 +59,14 @@ const numberOrNull = (value: string | number | null | undefined): number | null export const teamAdminSettingsChanges = ( values: TeamAdminSettingsValues, + initialValues: TeamAdminSettingsValues, editableFields: ReadonlySet, -): TeamAdminSettingsChanges => (editableFields.has("tpm_limit") ? { tpm_limit: numberOrNull(values.tpm_limit) } : {}); +): TeamAdminSettingsChanges => { + const tpmLimit = numberOrNull(values.tpm_limit); + return editableFields.has("tpm_limit") && tpmLimit !== numberOrNull(initialValues.tpm_limit) + ? { tpm_limit: tpmLimit } + : {}; +}; export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => { const parsed = callerEditAccessSchema.safeParse(callerEditAccess); From ada0a1ad3a06d3ae970a0e5e61d224867fb2fcf5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:49:29 -0700 Subject: [PATCH 89/89] fix(azure_ai): strip the azure_ai/ prefix when a Responses call is remapped to azure A catalog OpenAI name on an .openai.azure.com host (or with AZURE_AI_API_BASE set to one) is remapped from azure_ai to azure before the Responses request is built, and the azure_ai/ prefix stayed in the wire model, so Azure answered DeploymentNotFound. The Azure Responses config now strips azure_ai/ next to responses/ and o_series/. --- litellm/llms/azure/responses/transformation.py | 7 +------ .../response/test_azure_transformation.py | 11 +++++++++++ .../test_azure_ai_responses_transformation.py | 18 ++++++++++++++++++ 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 7fe12138ebc..2a82b42df7b 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) def get_stripped_model_name(self, model: str) -> str: - # if "responses/" is in the model name, remove it - if "responses/" in model: - model = model.replace("responses/", "") - if "o_series" in model: - model = model.replace("o_series/", "") - return model + return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "") def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]: """ diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 726c9f65681..532c278e891 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -677,3 +677,14 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo model="gpt-6-astra", drop_params=False, ) + + +def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap(): + request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request( + model="azure_ai/gpt-5.4-nano", + input="hi", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["model"] == "gpt-5.4-nano" diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 925608a3c9b..bae956eb061 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -253,6 +253,24 @@ async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, ap _assert_native_responses_request(route, expected_url, expected_model) +@pytest.mark.asyncio +@respx.mock +async def test_aresponses_catalog_name_remapped_to_azure_sends_bare_deployment_name(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://res.openai.azure.com") + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload("gpt-5.4-nano")) + ) + + await litellm.aresponses( + model="azure_ai/gpt-5.4-nano", + input="What is the weather in SF?", + api_base="https://res.openai.azure.com", + api_key="fake-key", + ) + + assert json.loads(route.calls.last.request.content)["model"] == "gpt-5.4-nano" + + @pytest.mark.asyncio @respx.mock @pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES)