From 518c4b07a72e989ed777f93c01ef296c77a8b567 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 21:55:44 +0000 Subject: [PATCH 01/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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 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 16/58] 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 17/58] 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 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 18/58] 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 19/58] 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 20/58] 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 21/58] 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 22/58] 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 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 23/58] 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 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 24/58] 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 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 25/58] 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 26/58] 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 b96a80400441aa4073c837e21d9c542a0aa3814e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 20:55:32 +0000 Subject: [PATCH 27/58] 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 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 28/58] 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 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 29/58] 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 9b77b5c2cb5ec8bbf58738d279e266e4f8162212 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 18:08:09 -0700 Subject: [PATCH 30/58] 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 37/58] 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 38/58] 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 39/58] 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 40/58] 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 41/58] 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 42/58] 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 43/58] 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 36eb9cdb356e0ad02124e7ae149324def11b4669 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:56:27 -0700 Subject: [PATCH 44/58] 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 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 45/58] 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 46/58] 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 47/58] 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 48/58] 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 49/58] 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 50/58] 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 51/58] 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 52/58] 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 53/58] 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 54/58] 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 55/58] 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 56/58] 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 57/58] 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 58/58] 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)