From ebb0f7e4cf1bfde5f720d89b76c4311176851878 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 18:40:32 +0000 Subject: [PATCH 001/130] fix(responses): preserve reasoning through prompt hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/main.py | 18 +++- litellm/responses/utils.py | 50 ++++++++++ .../test_responses_prompt_management.py | 94 ++++++++++++++++--- 3 files changed, 148 insertions(+), 14 deletions(-) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 12f9be970c7..453676937fc 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -494,7 +494,14 @@ async def aresponses( prompt_label=kwargs.get("prompt_label", None), prompt_version=kwargs.get("prompt_version", None), ) - input = cast(Union[str, ResponseInputParam], merged_input) + input = cast( + Union[str, ResponseInputParam], + ResponsesAPIRequestUtils.merge_prompt_management_input( + original_input=input, + client_input=client_input, + merged_input=merged_input, + ), + ) if model != original_model: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) kwargs.pop("prompt_id", None) @@ -609,7 +616,14 @@ def _apply_prompt_management_to_responses_call( prompt_label=kwargs.get("prompt_label", None), prompt_version=kwargs.get("prompt_version", None), ) - input = cast(Union[str, ResponseInputParam], merged_input) + input = cast( + Union[str, ResponseInputParam], + ResponsesAPIRequestUtils.merge_prompt_management_input( + original_input=input, + client_input=client_input, + merged_input=merged_input, + ), + ) local_vars["input"] = input local_vars["model"] = model if model != original_model: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 234eb777aca..6d42e33a268 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -19,7 +19,9 @@ import litellm from litellm._logging import verbose_logger from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.types.llms.openai import ( + AllMessageValues, ResponseAPIUsage, + ResponseInputParam, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponseText, @@ -36,6 +38,54 @@ from litellm.types.utils import ( class ResponsesAPIRequestUtils: """Helper utils for constructing ResponseAPI requests""" + @staticmethod + def merge_prompt_management_input( + original_input: str | ResponseInputParam, + client_input: list[AllMessageValues], + merged_input: list[AllMessageValues], + ) -> list[object]: + if isinstance(original_input, str): + return [*merged_input] + + original_items = tuple(original_input) + client_item_ids = frozenset(id(item) for item in client_input) + message_positions = tuple(index for index, item in enumerate(original_items) if id(item) in client_item_ids) + + if len(message_positions) == len(original_items): + return [*merged_input] + if not message_positions: + return [*merged_input, *original_items] + + corresponding_messages = len(client_input) == len(merged_input) and all( + original.get("role") == merged.get("role") + and (not isinstance(original.get("id"), str) or original.get("id") == merged.get("id")) + for original, merged in zip(client_input, merged_input) + ) + if corresponding_messages: + merged_by_position = dict(zip(message_positions, merged_input)) + return [ + merged_by_position[index] if index in merged_by_position else item + for index, item in enumerate(original_items) + ] + + all_messages_preserved = all(any(original is merged for merged in merged_input) for original in client_input) + if all_messages_preserved: + prefixes = { + id(original_items[position]): original_items[ + message_positions[index - 1] + 1 if index else 0 : position + ] + for index, position in enumerate(message_positions) + } + trailing_items = original_items[message_positions[-1] + 1 :] + return [item for merged in merged_input for item in (*prefixes.get(id(merged), ()), merged)] + list( + trailing_items + ) + + verbose_logger.warning( + "Prompt management hook replaced Responses API messages; non-message input items were dropped" + ) + return [*merged_input] + @staticmethod def _check_valid_arg( supported_params: Optional[List[str]], diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index 84e98390268..e4207b292da 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -14,13 +14,19 @@ Covers: """ import asyncio -from typing import List +from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import ( + AllMessageValues, + ResponseInputParam, +) # --------------------------------------------------------------------------- # Helpers @@ -54,18 +60,15 @@ def _patch_responses_dispatch(): return_value=("gpt-4o", "openai", None, None), ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler." - "LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", return_value=False, ), patch( - "litellm.responses.main.ProviderConfigManager" - ".get_provider_responses_api_config", + "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", return_value=None, ), patch( - "litellm.responses.main.litellm_completion_transformation_handler" - ".response_api_handler", + "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler", return_value=MagicMock(), ), ] @@ -77,7 +80,6 @@ def _patch_responses_dispatch(): class TestResponsesAPIPromptManagement: - def test_str_input_coerced_and_merged(self): """[A] str input is wrapped into a message list before being passed to the hook.""" template_messages: List[AllMessageValues] = [ @@ -108,9 +110,7 @@ class TestResponsesAPIPromptManagement: logging_obj.get_chat_completion_prompt.assert_called_once() call_kwargs = logging_obj.get_chat_completion_prompt.call_args.kwargs # str was coerced to a single user message before being passed to the hook - assert call_kwargs["messages"] == [ - {"role": "user", "content": "Tell me about AI."} - ] + assert call_kwargs["messages"] == [{"role": "user", "content": "Tell me about AI."}] assert call_kwargs["prompt_id"] == "summariser-prompt" def test_list_input_merged_with_template(self): @@ -256,6 +256,76 @@ class TestResponsesAPIPromptManagement: assert all(isinstance(m, dict) and "role" in m for m in passed_messages) assert len(passed_messages) == 1 + def test_cache_control_hook_preserves_reasoning_items(self): + system_message = cast( + AllMessageValues, + {"role": "system", "content": "Analyze the request"}, + ) + assistant_message = cast( + AllMessageValues, + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "The code has a bug", + "annotations": [], + } + ], + }, + ) + user_message = cast( + AllMessageValues, + {"role": "user", "content": "Check for security issues"}, + ) + reasoning_item = { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "encrypted", + } + original_input = cast( + ResponseInputParam, + [system_message, reasoning_item, assistant_message, user_message], + ) + _, merged_messages, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model="azure/gpt-5-codex", + messages=[system_message, assistant_message, user_message], + non_default_params={"cache_control_injection_points": [{"location": "message", "role": "system"}]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + logging_obj = _make_logging_obj( + merged_model="azure/gpt-5-codex", + merged_messages=merged_messages, + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + litellm.responses( + input=original_input, + model="azure/gpt-5-codex", + litellm_logging_obj=logging_obj, + cache_control_injection_points=[{"location": "message", "role": "system"}], + ) + + sent_input = mock_handler.call_args.kwargs["input"] + assert [item.get("type") for item in sent_input] == [ + None, + "reasoning", + "message", + None, + ] + assert sent_input[0]["cache_control"] == {"type": "ephemeral"} + assert sent_input[1] == reasoning_item + assert sent_input[2]["id"] == "msg_1" + def test_model_override_re_resolves_provider(self): """[G] When the prompt template overrides the model to a different provider, custom_llm_provider is re-resolved so downstream routing uses the correct provider. From 4baee71bdd1e82db5225122ad5d9a1a7ae925af2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 18:41:02 +0000 Subject: [PATCH 002/130] chore(responses): minimize regression test diff Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/test_responses_prompt_management.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index e4207b292da..b3ba81ee2e8 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -60,15 +60,18 @@ def _patch_responses_dispatch(): return_value=("gpt-4o", "openai", None, None), ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", + "litellm.responses.mcp.litellm_proxy_mcp_handler." + "LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", return_value=False, ), patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", + "litellm.responses.main.ProviderConfigManager" + ".get_provider_responses_api_config", return_value=None, ), patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler", + "litellm.responses.main.litellm_completion_transformation_handler" + ".response_api_handler", return_value=MagicMock(), ), ] @@ -80,6 +83,7 @@ def _patch_responses_dispatch(): class TestResponsesAPIPromptManagement: + def test_str_input_coerced_and_merged(self): """[A] str input is wrapped into a message list before being passed to the hook.""" template_messages: List[AllMessageValues] = [ @@ -110,7 +114,9 @@ class TestResponsesAPIPromptManagement: logging_obj.get_chat_completion_prompt.assert_called_once() call_kwargs = logging_obj.get_chat_completion_prompt.call_args.kwargs # str was coerced to a single user message before being passed to the hook - assert call_kwargs["messages"] == [{"role": "user", "content": "Tell me about AI."}] + assert call_kwargs["messages"] == [ + {"role": "user", "content": "Tell me about AI."} + ] assert call_kwargs["prompt_id"] == "summariser-prompt" def test_list_input_merged_with_template(self): From 4db0bdf465c9f07e8561136cf055ca9e6854f086 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 18:51:07 +0000 Subject: [PATCH 003/130] fix(responses): handle non-message-only prompt input Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/utils.py | 5 +- .../test_responses_prompt_management.py | 154 +++++++++++++----- 2 files changed, 116 insertions(+), 43 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 6d42e33a268..a5203c4ee6a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -54,7 +54,10 @@ class ResponsesAPIRequestUtils: if len(message_positions) == len(original_items): return [*merged_input] if not message_positions: - return [*merged_input, *original_items] + verbose_logger.warning( + "Prompt management hook returned messages without Responses API input messages; merged messages were ignored" + ) + return [*original_items] corresponding_messages = len(client_input) == len(merged_input) and all( original.get("role") == merged.get("role") diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index b3ba81ee2e8..7044d8384f8 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -77,6 +77,56 @@ def _patch_responses_dispatch(): ] +def _make_cache_control_case() -> tuple[ + ResponseInputParam, + list[AllMessageValues], + dict[str, object], +]: + system_message = cast( + AllMessageValues, + {"role": "system", "content": "Analyze the request"}, + ) + assistant_message = cast( + AllMessageValues, + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "The code has a bug", + "annotations": [], + } + ], + }, + ) + user_message = cast( + AllMessageValues, + {"role": "user", "content": "Check for security issues"}, + ) + reasoning_item = { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "encrypted", + } + original_input = cast( + ResponseInputParam, + [system_message, reasoning_item, assistant_message, user_message], + ) + _, merged_messages, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model="azure/gpt-5-codex", + messages=[system_message, assistant_message, user_message], + non_default_params={"cache_control_injection_points": [{"location": "message", "role": "system"}]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return original_input, merged_messages, reasoning_item + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -263,48 +313,7 @@ class TestResponsesAPIPromptManagement: assert len(passed_messages) == 1 def test_cache_control_hook_preserves_reasoning_items(self): - system_message = cast( - AllMessageValues, - {"role": "system", "content": "Analyze the request"}, - ) - assistant_message = cast( - AllMessageValues, - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "The code has a bug", - "annotations": [], - } - ], - }, - ) - user_message = cast( - AllMessageValues, - {"role": "user", "content": "Check for security issues"}, - ) - reasoning_item = { - "type": "reasoning", - "id": "rs_1", - "summary": [], - "encrypted_content": "encrypted", - } - original_input = cast( - ResponseInputParam, - [system_message, reasoning_item, assistant_message, user_message], - ) - _, merged_messages, _ = AnthropicCacheControlHook().get_chat_completion_prompt( - model="azure/gpt-5-codex", - messages=[system_message, assistant_message, user_message], - non_default_params={"cache_control_injection_points": [{"location": "message", "role": "system"}]}, - prompt_id=None, - prompt_variables=None, - dynamic_callback_params={}, - ) + original_input, merged_messages, reasoning_item = _make_cache_control_case() logging_obj = _make_logging_obj( merged_model="azure/gpt-5-codex", merged_messages=merged_messages, @@ -332,6 +341,37 @@ class TestResponsesAPIPromptManagement: assert sent_input[1] == reasoning_item assert sent_input[2]["id"] == "msg_1" + def test_all_non_message_input_items_remain_unchanged(self): + reasoning_item = { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "encrypted", + } + original_input = cast(ResponseInputParam, [reasoning_item]) + logging_obj = _make_logging_obj( + merged_model="openai/gpt-4o", + merged_messages=[ + cast( + AllMessageValues, + {"role": "system", "content": "Analyze the request"}, + ) + ], + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + litellm.responses( + input=original_input, + model="gpt-4o", + prompt_id="all-non-message", + litellm_logging_obj=logging_obj, + ) + + assert mock_handler.call_args.kwargs["input"] == original_input + def test_model_override_re_resolves_provider(self): """[G] When the prompt template overrides the model to a different provider, custom_llm_provider is re-resolved so downstream routing uses the correct provider. @@ -469,3 +509,33 @@ class TestAsyncResponsesAPIPromptManagement: passed_messages = call_kwargs["messages"] assert all(isinstance(m, dict) and "role" in m for m in passed_messages) assert len(passed_messages) == 1 + + @pytest.mark.asyncio + async def test_async_cache_control_hook_preserves_reasoning_items(self): + original_input, merged_messages, reasoning_item = _make_cache_control_case() + logging_obj = _make_logging_obj( + merged_model="azure/gpt-5-codex", + merged_messages=merged_messages, + ) + + patches = _patch_responses_dispatch() + with patches[0], patches[1], patches[2], patches[3] as mock_handler: + import litellm + + await litellm.aresponses( + input=original_input, + model="azure/gpt-5-codex", + litellm_logging_obj=logging_obj, + cache_control_injection_points=[{"location": "message", "role": "system"}], + ) + + sent_input = mock_handler.call_args.kwargs["input"] + assert [item.get("type") for item in sent_input] == [ + None, + "reasoning", + "message", + None, + ] + assert sent_input[0]["cache_control"] == {"type": "ephemeral"} + assert sent_input[1] == reasoning_item + assert sent_input[2]["id"] == "msg_1" From 5e68a003476d0a2ecce31036ea25597cf0a549d7 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 21 Jul 2026 00:10:49 +0000 Subject: [PATCH 004/130] test(e2e): add live A2A agent e2e suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 1 + tests/e2e/a2a/a2a_client.py | 217 +++++++++++++++++++++++++ tests/e2e/a2a/conftest.py | 17 ++ tests/e2e/a2a/test_a2a_agent_e2e.py | 139 ++++++++++++++++ tests/e2e/coverage_registry/other.yaml | 6 + 5 files changed, 380 insertions(+) create mode 100644 tests/e2e/a2a/a2a_client.py create mode 100644 tests/e2e/a2a/conftest.py create mode 100644 tests/e2e/a2a/test_a2a_agent_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 47f3c74d7f1..c35fb2fa435 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,6 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright) +- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server only (see "MCP suite: real Datadog only" below) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py new file mode 100644 index 00000000000..5274ec15383 --- /dev/null +++ b/tests/e2e/a2a/a2a_client.py @@ -0,0 +1,217 @@ +"""Client for the proxy's A2A (agent-to-agent) surface. + +An A2A agent is registered admin-side via POST /v1/agents with an agent card and +litellm_params; the proxy fronts it at /a2a/{id}, serving a proxy-owned agent card +at /.well-known/agent-card.json and accepting A2A JSON-RPC calls at /a2a/{id}. This +suite registers agents backed by the litellm_completion_bridge (custom_llm_provider ++ model), so message/send runs a real provider completion and comes back in the +agent's pinned A2A protocol version. The A2A request/response models are co-located +here because only this suite uses them. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass + +from pydantic import BaseModel, ConfigDict, Field + +from e2e_http import NoBody, Result, is_ok +from proxy_client import ProxyClient + + +class A2ACapabilities(BaseModel): + streaming: bool | None = None + push_notifications: bool | None = Field(default=None, serialization_alias="pushNotifications") + + +class A2ASkill(BaseModel): + id: str + name: str + description: str + tags: list[str] + + +class AgentCardParams(BaseModel): + """The upstream agent card an admin registers. `protocolVersion` is the field the + proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration.""" + + protocol_version: str = Field(serialization_alias="protocolVersion") + name: str + description: str + version: str + capabilities: A2ACapabilities = A2ACapabilities() + skills: list[A2ASkill] + default_input_modes: list[str] = Field(default=["text"], serialization_alias="defaultInputModes") + default_output_modes: list[str] = Field(default=["text"], serialization_alias="defaultOutputModes") + + +class A2ABridgeParams(BaseModel): + """litellm_params that route the agent through the completion bridge: an A2A + message/send is transformed into a litellm.acompletion against this provider.""" + + model_config = ConfigDict(protected_namespaces=()) + + custom_llm_provider: str + model: str + + +class AgentRegisterBody(BaseModel): + agent_name: str + agent_card_params: AgentCardParams + litellm_params: A2ABridgeParams + + +class A2ASecurityScheme(BaseModel): + type: str + scheme: str + + +class A2AInterface(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + url: str + protocol_version: str | None = Field(default=None, alias="protocolVersion") + + +class ServedAgentCard(BaseModel): + """The proxy-owned card, either nested under a registration response's + `agent_card_params` or served raw at /.well-known/agent-card.json. The proxy + rewrites `url`/`supportedInterfaces` to itself and replaces the security scheme + with its own virtual-key bearer scheme.""" + + model_config = ConfigDict(populate_by_name=True) + + protocol_version: str = Field(alias="protocolVersion") + name: str + url: str | None = None + security_schemes: dict[str, A2ASecurityScheme] | None = Field(default=None, alias="securitySchemes") + security: list[dict[str, list[str]]] | None = None + supported_interfaces: list[A2AInterface] | None = Field(default=None, alias="supportedInterfaces") + + +class AgentResponse(BaseModel): + agent_id: str + agent_name: str + agent_card_params: ServedAgentCard + + +class A2ATextPart(BaseModel): + kind: str = "text" + text: str + + +class A2AOutboundMessage(BaseModel): + role: str = "user" + parts: list[A2ATextPart] + message_id: str = Field(serialization_alias="messageId") + + +class A2AMessageSendParams(BaseModel): + message: A2AOutboundMessage + + +class A2AJsonRpcRequest(BaseModel): + jsonrpc: str = "2.0" + id: str + method: str = "message/send" + params: A2AMessageSendParams + + +class A2AResponsePart(BaseModel): + kind: str | None = None + text: str | None = None + + +class A2AResponseMessage(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + message_id: str | None = Field(default=None, alias="messageId") + role: str | None = None + parts: list[A2AResponsePart] = [] + + +class A2AResult(BaseModel): + """A message/send result. In 0.3 the message fields sit directly on the result + (`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the + agent's reply from whichever shape the served version produced.""" + + model_config = ConfigDict(populate_by_name=True) + + kind: str | None = None + role: str | None = None + message_id: str | None = Field(default=None, alias="messageId") + parts: list[A2AResponsePart] = [] + message: A2AResponseMessage | None = None + + @property + def text(self) -> str: + parts = self.message.parts if self.message is not None else self.parts + return "".join(part.text or "" for part in parts) + + @property + def is_nested_v1_shape(self) -> bool: + return self.message is not None + + +class A2AError(BaseModel): + code: int + message: str + + +class A2AResponse(BaseModel): + jsonrpc: str + id: str | None = None + result: A2AResult | None = None + error: A2AError | None = None + + +@dataclass(frozen=True, slots=True) +class A2AClient: + proxy: ProxyClient + + def register_agent(self, body: AgentRegisterBody) -> Result[AgentResponse]: + return self.proxy.transport.post( + "/v1/agents", + headers=self.proxy.transport.master, + json=body, + response_type=AgentResponse, + ) + + def get_agent(self, agent_id: str) -> Result[AgentResponse]: + return self.proxy.transport.get( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AgentResponse, + ) + + def delete_agent(self, agent_id: str) -> None: + result = self.proxy.transport.delete( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_agent({agent_id!r}) failed: {result}", stacklevel=2) + + def agent_card(self, agent_id: str, key: str) -> Result[ServedAgentCard]: + return self.proxy.transport.get( + f"/a2a/{agent_id}/.well-known/agent-card.json", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=ServedAgentCard, + ) + + def send_message(self, agent_id: str, key: str, body: A2AJsonRpcRequest) -> Result[A2AResponse]: + return self.proxy.transport.post( + f"/a2a/{agent_id}", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=A2AResponse, + ) + + +def build_a2a_client(proxy: ProxyClient) -> A2AClient: + return A2AClient(proxy=proxy) diff --git a/tests/e2e/a2a/conftest.py b/tests/e2e/a2a/conftest.py new file mode 100644 index 00000000000..93f3b56c8f7 --- /dev/null +++ b/tests/e2e/a2a/conftest.py @@ -0,0 +1,17 @@ +"""A2A suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker +live in the parent tests/e2e/conftest.py. A2AClient holds the shared ProxyClient, +so the `resources` fixture cleans up keys this suite creates; agents are torn down +via `resources.defer(...)` in each test. +""" + +import pytest + +from a2a_client import A2AClient, build_a2a_client +from proxy_client import ProxyClient + + +@pytest.fixture(scope="session") +def client(proxy: ProxyClient) -> A2AClient: + return build_a2a_client(proxy) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py new file mode 100644 index 00000000000..eb61ace238c --- /dev/null +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -0,0 +1,139 @@ +"""A2A agents end to end, against a live proxy. + +An admin registers an agent whose card pins an A2A protocol version and whose +litellm_params route it through the completion bridge; a caller then discovers the +proxy-owned card and drives it over A2A JSON-RPC. These tests assert the recorded +state (the agent persists, a spend row lands) and the enforced behavior (the served +card points back at the proxy, message/send returns a real completion in the pinned +protocol version, and an unsupported version is refused at registration). +""" + +from __future__ import annotations + +import pytest + +from a2a_client import ( + A2ABridgeParams, + A2AClient, + A2AJsonRpcRequest, + A2AMessageSendParams, + A2AOutboundMessage, + A2ASkill, + A2ATextPart, + AgentCardParams, + AgentRegisterBody, + AgentResponse, +) +from e2e_config import unique_marker +from e2e_http import UnknownApiError, unwrap +from lifecycle import ResourceManager + +BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5") + +pytestmark = pytest.mark.e2e + + +def _register(client: A2AClient, resources: ResourceManager, protocol_version: str) -> AgentResponse: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-{marker}", + agent_card_params=AgentCardParams( + protocol_version=protocol_version, + name=f"E2E A2A {marker}", + description="e2e agent backed by the litellm completion bridge", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="general chat", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + agent = unwrap(client.register_agent(body)) + resources.defer(lambda: client.delete_agent(agent.agent_id)) + return agent + + +def _ask(text: str) -> A2AJsonRpcRequest: + return A2AJsonRpcRequest( + id=f"e2e-{unique_marker()}", + params=A2AMessageSendParams( + message=A2AOutboundMessage(parts=[A2ATextPart(text=text)], message_id=unique_marker()) + ), + ) + + +class TestA2AAgentLifecycle: + @pytest.mark.covers("other.a2a.register.persists") + def test_register_persists(self, client: A2AClient, resources: ResourceManager) -> None: + agent = _register(client, resources, "0.3") + fetched = unwrap(client.get_agent(agent.agent_id)) + assert fetched.agent_id == agent.agent_id + assert fetched.agent_name == agent.agent_name + assert fetched.agent_card_params.protocol_version == "0.3" + + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") + def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.url is not None and card.url.endswith(f"/a2a/{agent.agent_id}") + assert card.security_schemes is not None + scheme = next(iter(card.security_schemes.values())) + assert scheme.scheme == "bearer" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].url == card.url + + @pytest.mark.covers("other.a2a.message_send.bridge_invokes") + def test_message_send_runs_completion_bridge(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Reply with exactly the word PONG and nothing else") + response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) + assert response.error is None + assert response.result is not None + assert "PONG" in response.result.text.upper() + + rows = client.proxy.poll_logs_for_request_id(request.id) + assert rows, f"no spend log row landed for a2a request {request.id}" + assert rows[0].call_type == "asend_message" + assert rows[0].model == f"a2a_agent/{agent.agent_card_params.name}" + + @pytest.mark.covers("other.a2a.version.serves_pinned_0_3") + def test_pinned_v0_3_serves_flat_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert not result.is_nested_v1_shape + assert result.kind == "message" + assert result.role == "agent" + assert result.text != "" + + @pytest.mark.covers("other.a2a.version.serves_pinned_1_0") + def test_pinned_v1_0_serves_nested_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "1.0") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert result.is_nested_v1_shape + assert result.message is not None + assert result.message.role == "ROLE_AGENT" + assert result.text != "" + + @pytest.mark.covers("other.a2a.register.unsupported_version_rejected") + def test_unsupported_protocol_version_rejected(self, client: A2AClient, resources: ResourceManager) -> None: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-bad-{marker}", + agent_card_params=AgentCardParams( + protocol_version="9.9", + name=f"E2E A2A bad {marker}", + description="unsupported version", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="c", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + result = client.register_agent(body) + match result: + case UnknownApiError(status_code=status, body=detail): + assert status == 400 + assert "protocolVersion" in detail + case _: + pytest.fail(f"expected 400 for unsupported protocolVersion, got {result}") diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 6b183cbf9f3..f4d0120e085 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -28,3 +28,9 @@ - {id: other.config.overrides.audit_logged, module: other, tier: P1, area: config, assertions: [audit_logged], source: "config_override_endpoints.py:67-100", rationale: "Config override mutations audit-logged, values redacted"} - {id: other.key_mgmt.regenerate.grace_period_honored, module: other, tier: P1, area: auth, assertions: [grace_period_honored], source: "key_management_endpoints.py:4503-4560", rationale: "Old key valid during grace_period then revoked"} - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} +- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} +- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} +- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} +- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} +- {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} +- {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} From 1315ebd1f97c2c3bb56f278a45d1904d48657401 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 21 Jul 2026 20:14:21 +0000 Subject: [PATCH 005/130] test(e2e): guard 0.3.0-style semver protocolVersion registration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/a2a/test_a2a_agent_e2e.py | 10 ++++++++++ tests/e2e/coverage_registry/other.yaml | 1 + 2 files changed, 11 insertions(+) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index eb61ace238c..823f6b9c001 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -69,6 +69,16 @@ class TestA2AAgentLifecycle: assert fetched.agent_name == agent.agent_name assert fetched.agent_card_params.protocol_version == "0.3" + @pytest.mark.covers("other.a2a.register.semver_version_accepted") + def test_semver_protocol_version_registers_and_serves(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3.0") + assert agent.agent_card_params.protocol_version.startswith("0.3") + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.protocol_version.startswith("0.3") + result = unwrap(client.send_message(agent.agent_id, scoped_key, _ask("Say hi in one word"))).result + assert result is not None + assert result.text != "" + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: agent = _register(client, resources, "0.3") diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index f4d0120e085..63a626caab4 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -30,6 +30,7 @@ - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} - {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} - {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} +- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers and serves as the 0.3 family rather than 400ing; regression guard for the v1.92 report"} - {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} From 43402b04bd1914d10fe0d011f666791e2c68888f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:43:09 -0700 Subject: [PATCH 006/130] test(e2e): register the verbatim a2a-sdk 0.3.x card and guard malformed protocolVersion --- tests/e2e/a2a/a2a_client.py | 3 + tests/e2e/a2a/test_a2a_agent_e2e.py | 95 +++++++++++++++++++++----- tests/e2e/coverage_registry/other.yaml | 4 +- 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index 5274ec15383..c1a1503ec96 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -30,6 +30,7 @@ class A2ASkill(BaseModel): name: str description: str tags: list[str] + examples: list[str] | None = None class AgentCardParams(BaseModel): @@ -40,10 +41,12 @@ class AgentCardParams(BaseModel): name: str description: str version: str + url: str | None = None capabilities: A2ACapabilities = A2ACapabilities() skills: list[A2ASkill] default_input_modes: list[str] = Field(default=["text"], serialization_alias="defaultInputModes") default_output_modes: list[str] = Field(default=["text"], serialization_alias="defaultOutputModes") + preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport") class A2ABridgeParams(BaseModel): diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index 823f6b9c001..8c19f97e64c 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -14,6 +14,7 @@ import pytest from a2a_client import ( A2ABridgeParams, + A2ACapabilities, A2AClient, A2AJsonRpcRequest, A2AMessageSendParams, @@ -25,7 +26,7 @@ from a2a_client import ( AgentResponse, ) from e2e_config import unique_marker -from e2e_http import UnknownApiError, unwrap +from e2e_http import Result, UnknownApiError, unwrap from lifecycle import ResourceManager BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5") @@ -51,6 +52,50 @@ def _register(client: A2AClient, resources: ResourceManager, protocol_version: s return agent +def _google_sdk_default_card(marker: str) -> AgentCardParams: + """Field-for-field the card the Google a2a-sdk 0.3.x line serves for its helloworld + sample, whose ``AgentCard`` defaults ``protocolVersion`` to the full semver "0.3.0" + (the shape the v1.92 regression rejected); ``url`` and ``name`` are the + deployment-specific fields the SDK requires callers to fill. The url points at a + reserved example host: a url-bearing card makes message/send dial that upstream + rather than the completion bridge, so the verbatim-card test asserts registration + and card serving only.""" + return AgentCardParams( + protocol_version="0.3.0", + name=f"E2E A2A sdk {marker}", + description="Just a hello world agent", + version="1.0.0", + url="http://e2e-a2a-upstream.example/", + capabilities=A2ACapabilities(streaming=True), + skills=[ + A2ASkill( + id="hello_world", + name="Returns hello world", + description="just returns hello world", + tags=["hello world"], + examples=["hi", "hello world"], + ) + ], + preferred_transport="JSONRPC", + ) + + +def _register_rejection(client: A2AClient, protocol_version: str) -> Result[AgentResponse]: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-bad-{marker}", + agent_card_params=AgentCardParams( + protocol_version=protocol_version, + name=f"E2E A2A bad {marker}", + description="rejected at registration", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="c", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + return client.register_agent(body) + + def _ask(text: str) -> A2AJsonRpcRequest: return A2AJsonRpcRequest( id=f"e2e-{unique_marker()}", @@ -72,13 +117,31 @@ class TestA2AAgentLifecycle: @pytest.mark.covers("other.a2a.register.semver_version_accepted") def test_semver_protocol_version_registers_and_serves(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: agent = _register(client, resources, "0.3.0") - assert agent.agent_card_params.protocol_version.startswith("0.3") + assert agent.agent_card_params.protocol_version == "0.3" card = unwrap(client.agent_card(agent.agent_id, scoped_key)) - assert card.protocol_version.startswith("0.3") + assert card.protocol_version == "0.3" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].protocol_version == "0.3" result = unwrap(client.send_message(agent.agent_id, scoped_key, _ask("Say hi in one word"))).result assert result is not None assert result.text != "" + @pytest.mark.covers("other.a2a.register.sdk_default_card_accepted") + def test_google_sdk_default_card_registers_verbatim(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-sdk-{marker}", + agent_card_params=_google_sdk_default_card(marker), + litellm_params=BRIDGE, + ) + agent = unwrap(client.register_agent(body)) + resources.defer(lambda: client.delete_agent(agent.agent_id)) + assert agent.agent_card_params.protocol_version == "0.3" + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.protocol_version == "0.3" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].protocol_version == "0.3" + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: agent = _register(client, resources, "0.3") @@ -127,23 +190,21 @@ class TestA2AAgentLifecycle: assert result.text != "" @pytest.mark.covers("other.a2a.register.unsupported_version_rejected") - def test_unsupported_protocol_version_rejected(self, client: A2AClient, resources: ResourceManager) -> None: - marker = unique_marker() - body = AgentRegisterBody( - agent_name=f"e2e-a2a-bad-{marker}", - agent_card_params=AgentCardParams( - protocol_version="9.9", - name=f"E2E A2A bad {marker}", - description="unsupported version", - version="1.0.0", - skills=[A2ASkill(id="chat", name="Chat", description="c", tags=["chat"])], - ), - litellm_params=BRIDGE, - ) - result = client.register_agent(body) + def test_unsupported_protocol_version_rejected(self, client: A2AClient) -> None: + result = _register_rejection(client, "9.9") match result: case UnknownApiError(status_code=status, body=detail): assert status == 400 assert "protocolVersion" in detail case _: pytest.fail(f"expected 400 for unsupported protocolVersion, got {result}") + + @pytest.mark.covers("other.a2a.register.malformed_version_rejected") + def test_malformed_protocol_version_rejected(self, client: A2AClient) -> None: + result = _register_rejection(client, "0.3.garbage") + match result: + case UnknownApiError(status_code=status, body=detail): + assert status == 400 + assert "Unsupported protocolVersion '0.3.garbage'" in detail + case _: + pytest.fail(f"expected 400 for malformed protocolVersion, got {result}") diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 63a626caab4..00a76cf511a 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -30,7 +30,9 @@ - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} - {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} - {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} -- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers and serves as the 0.3 family rather than 400ing; regression guard for the v1.92 report"} +- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"} +- {id: other.a2a.register.sdk_default_card_accepted, module: other, tier: P1, area: a2a, assertions: [sdk_default_card_accepted], source: "agent_endpoints/endpoints.py _build_merged_agent_card", rationale: "The full field set the Google a2a-sdk 0.3.x emits (protocolVersion 0.3.0, url, preferredTransport, capabilities, skill examples) registers verbatim and serves the canonical 0.3"} +- {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"} - {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} From b572cb80d1058cfea4ec01b8f9de36fa8f1f2702 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:53:23 -0700 Subject: [PATCH 007/130] test(e2e): add weekly session-anomaly load test against real providers --- .github/workflows/weekly_load_anomaly.yml | 81 ++++++ tests/e2e/CLAUDE.md | 4 +- tests/e2e/conftest.py | 4 + tests/e2e/coverage_registry/reliability.yaml | 1 + tests/e2e/e2e_config.py | 14 + tests/e2e/load/conftest.py | 18 ++ tests/e2e/load/session_anomaly.py | 239 ++++++++++++++++++ .../load/test_weekly_session_anomaly_e2e.py | 131 ++++++++++ tests/e2e/load/weekly_anomaly_config.yml | 3 + tests/e2e/pytest.ini | 1 + 10 files changed, 494 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/weekly_load_anomaly.yml create mode 100644 tests/e2e/load/session_anomaly.py create mode 100644 tests/e2e/load/test_weekly_session_anomaly_e2e.py create mode 100644 tests/e2e/load/weekly_anomaly_config.yml diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml new file mode 100644 index 00000000000..6853ffa1cde --- /dev/null +++ b/.github/workflows/weekly_load_anomaly.yml @@ -0,0 +1,81 @@ +name: "Weekly Load Anomaly Check" + +on: + schedule: + - cron: "0 6 * * 1" + workflow_dispatch: + +permissions: + contents: read + +jobs: + weekly-load-anomaly: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 45 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-weekly-anomaly-check + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Generate Prisma client + env: + PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy + run: | + nohup uv run --no-sync litellm --config tests/e2e/load/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the weekly session anomaly test + env: + E2E_WEEKLY_ANOMALY: "1" + run: | + uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index bf34a896771..daf9fa9e81b 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) -- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites +- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites. Also home of the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`): Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; additionally marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, because it spends real provider money (driven by `.github/workflows/weekly_load_anomaly.yml`) - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke @@ -131,7 +131,7 @@ reliability... behavior : fallback | retry | cooldown | timeout | routing | cache | circuit_breaker | perf variant : 5xx | context_window | content_policy | 429 | timeout simple_shuffle | usage_based | latency_based | cost_based | least_busy - latency | throughput (perf only; SLO/threshold assertion, not binary) + latency | throughput | session_anomaly (perf only; SLO/threshold assertion, not binary) assertion : routes_to_fallback | succeeds_within_retries | picks_under_tpm | returns_cached | trips_then_recovers | under_slo e.g. reliability.fallback.context_window.routes_to_fallback exercised_on=[chat_completions] diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 609da6a9b07..7b5ad3551b8 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -43,6 +43,10 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites", ) + config.addinivalue_line( + "markers", + "weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set", + ) def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 1538d3f3cda..ebbfd3415a5 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -24,3 +24,4 @@ - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} +- {id: reliability.perf.session_anomaly.under_slo, module: reliability, tier: P1, behavior: perf, variant: session_anomaly, assertions: [under_slo], exercised_on: [messages], source: grammar, rationale: "Weekly Claude Code-shaped multi-turn session load against real providers; ceilings on error rate, warm-turn cache read/write, p95 turn time, and gateway-recorded spend (LIT-4562)"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 3be339d28a0..985b299eac8 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -78,6 +78,20 @@ LOAD_DURATION_SECONDS = float(os.environ.get("E2E_LOAD_DURATION_SECONDS", "60")) LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355")) LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01")) +WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" +ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) +ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) +ANOMALY_MAX_ERROR_RATIO = float(os.environ.get("E2E_ANOMALY_MAX_ERROR_RATIO", "0.05")) +ANOMALY_MIN_WARM_CACHE_READ_SHARE = float( + os.environ.get("E2E_ANOMALY_MIN_WARM_CACHE_READ_SHARE", "0.5") +) +ANOMALY_MAX_P95_TURN_SECONDS = float( + os.environ.get("E2E_ANOMALY_MAX_P95_TURN_SECONDS", "30") +) +ANOMALY_MAX_KEY_SPEND_USD = float( + os.environ.get("E2E_ANOMALY_MAX_KEY_SPEND_USD", "0.60") +) + def require_env(*names: str) -> tuple[str, ...]: """Return the non-empty values for each env name, or hard-fail naming which are missing. diff --git a/tests/e2e/load/conftest.py b/tests/e2e/load/conftest.py index e9fba02680d..d9b1b2d9d69 100644 --- a/tests/e2e/load/conftest.py +++ b/tests/e2e/load/conftest.py @@ -1,10 +1,12 @@ from __future__ import annotations +import os from collections.abc import Iterator import pytest from requests import RequestException +from e2e_config import WEEKLY_ANOMALY_OPT_IN_ENV from e2e_http import NoBody, Success from load_client import LoadClient, build_client from load_constants import LOAD_MODEL @@ -18,6 +20,22 @@ LOAD_MODEL_PARAMS = LiteLLMParamsBody( ) +def pytest_collection_modifyitems( + config: pytest.Config, items: list[pytest.Item] +) -> None: + if os.environ.get(WEEKLY_ANOMALY_OPT_IN_ENV): + return + deselected = [ + item for item in items if item.get_closest_marker("weekly") is not None + ] + if not deselected: + return + config.hook.pytest_deselected(items=deselected) + items[:] = [ + item for item in items if item.get_closest_marker("weekly") is None + ] + + @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> LoadClient: return build_client(proxy) diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/load/session_anomaly.py new file mode 100644 index 00000000000..fc1cfa7ac70 --- /dev/null +++ b/tests/e2e/load/session_anomaly.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import Result, Success +from models import CacheControl, RichMessage, TextBlock +from transport import Transport + + +class SessionMessagesRequest(BaseModel): + model: str + max_tokens: int = 128 + system: list[TextBlock] + messages: list[RichMessage] + + +class SessionUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class SessionContentBlock(BaseModel): + type: str | None = None + text: str | None = None + + +class SessionMessagesResponse(BaseModel): + content: list[SessionContentBlock] = [] + usage: SessionUsage = SessionUsage() + + @property + def text(self) -> str: + return "".join(block.text or "" for block in self.content) + + +@dataclass(frozen=True, slots=True) +class TurnMetric: + turn_index: int + ok: bool + latency_seconds: float + uncached_input_tokens: int + cache_read_tokens: int + cache_creation_tokens: int + failure: str | None + + +@dataclass(frozen=True, slots=True) +class AnomalyReport: + attempted_turns: int + failed_turns: int + warm_turns: int + warm_uncached_input_tokens: int + warm_cache_read_tokens: int + warm_cache_creation_tokens: int + p95_turn_seconds: float + + @property + def error_ratio(self) -> float: + return self.failed_turns / self.attempted_turns if self.attempted_turns else 1.0 + + @property + def warm_cache_read_share(self) -> float: + billed = ( + self.warm_uncached_input_tokens + + self.warm_cache_read_tokens + + self.warm_cache_creation_tokens + ) + return self.warm_cache_read_tokens / billed if billed else 0.0 + + +def _system_prefix_block(marker: str) -> TextBlock: + text = " ".join( + f"Project context paragraph {index} for session {marker}." for index in range(300) + ) + return TextBlock(text=text, cache_control=CacheControl()) + + +def _user_turn_text(marker: str, turn_index: int) -> str: + notes = " ".join( + f"Working note {index} of turn {turn_index} in session {marker}." + for index in range(80) + ) + return f"Reply with one short sentence.\n{notes}" + + +def _reminder_turn() -> RichMessage: + return RichMessage( + role="system", + content=[ + TextBlock( + text="Keep the answer to one short sentence." + ) + ], + ) + + +def _without_cache_control(message: RichMessage) -> RichMessage: + return RichMessage( + role=message.role, + content=[TextBlock(text=block.text) for block in message.content], + ) + + +def _metric( + result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float +) -> TurnMetric: + if isinstance(result, Success): + usage = result.data.usage + return TurnMetric( + turn_index=turn_index, + ok=True, + latency_seconds=latency_seconds, + uncached_input_tokens=usage.input_tokens, + cache_read_tokens=usage.cache_read_input_tokens, + cache_creation_tokens=usage.cache_creation_input_tokens, + failure=None, + ) + return TurnMetric( + turn_index=turn_index, + ok=False, + latency_seconds=latency_seconds, + uncached_input_tokens=0, + cache_read_tokens=0, + cache_creation_tokens=0, + failure=repr(result), + ) + + +def _drive_turns( + transport: Transport, + key: str, + model: str, + marker: str, + system_block: TextBlock, + history: tuple[RichMessage, ...], + turn_index: int, + remaining_turns: int, +) -> tuple[TurnMetric, ...]: + if remaining_turns == 0: + return () + user_turn = RichMessage( + role="user", + content=[ + TextBlock( + text=_user_turn_text(marker, turn_index), cache_control=CacheControl() + ) + ], + ) + started = time.monotonic() + result = transport.post( + "/v1/messages", + headers=transport.bearer(key), + json=SessionMessagesRequest( + model=model, + system=[system_block], + messages=[*history, user_turn], + ), + response_type=SessionMessagesResponse, + ) + turn = _metric(result, turn_index, time.monotonic() - started) + if not isinstance(result, Success): + return (turn,) + assistant_turn = RichMessage( + role="assistant", content=[TextBlock(text=result.data.text or "Understood.")] + ) + return ( + turn, + *_drive_turns( + transport, + key, + model, + marker, + system_block, + ( + *history, + _without_cache_control(user_turn), + _reminder_turn(), + assistant_turn, + ), + turn_index + 1, + remaining_turns - 1, + ), + ) + + +def run_session( + transport: Transport, key: str, model: str, turns: int +) -> tuple[TurnMetric, ...]: + marker = unique_marker() + return _drive_turns( + transport, + key, + model, + marker, + _system_prefix_block(marker), + (), + 1, + turns, + ) + + +def run_concurrent_sessions( + transport: Transport, key: str, model: str, sessions: int, turns_per_session: int +) -> tuple[TurnMetric, ...]: + with ThreadPoolExecutor(max_workers=sessions) as pool: + futures = [ + pool.submit(run_session, transport, key, model, turns_per_session) + for _ in range(sessions) + ] + return tuple(turn for future in futures for turn in future.result()) + + +def _p95(latencies: tuple[float, ...]) -> float: + if not latencies: + return 0.0 + ranked = sorted(latencies) + return ranked[max(0, -(-len(ranked) * 95 // 100) - 1)] + + +def summarize(turns: tuple[TurnMetric, ...]) -> AnomalyReport: + warm = tuple(turn for turn in turns if turn.ok and turn.turn_index >= 2) + return AnomalyReport( + attempted_turns=len(turns), + failed_turns=sum(1 for turn in turns if not turn.ok), + warm_turns=len(warm), + warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm), + warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm), + warm_cache_creation_tokens=sum(turn.cache_creation_tokens for turn in warm), + p95_turn_seconds=_p95( + tuple(turn.latency_seconds for turn in turns if turn.ok) + ), + ) diff --git a/tests/e2e/load/test_weekly_session_anomaly_e2e.py b/tests/e2e/load/test_weekly_session_anomaly_e2e.py new file mode 100644 index 00000000000..ebd58166f0b --- /dev/null +++ b/tests/e2e/load/test_weekly_session_anomaly_e2e.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import time +from dataclasses import dataclass + +import pytest + +from e2e_config import ( + ANOMALY_MAX_ERROR_RATIO, + ANOMALY_MAX_KEY_SPEND_USD, + ANOMALY_MAX_P95_TURN_SECONDS, + ANOMALY_MIN_WARM_CACHE_READ_SHARE, + ANOMALY_SESSIONS, + ANOMALY_TURNS_PER_SESSION, + unique_marker, +) +from lifecycle import ResourceManager +from load_client import LoadClient +from models import KeyGenerateBody, LiteLLMParamsBody +from proxy_client import ProxyClient +from session_anomaly import run_concurrent_sessions, summarize + +pytestmark = [pytest.mark.e2e, pytest.mark.load, pytest.mark.weekly] + + +@dataclass(frozen=True, slots=True) +class AnomalyRoute: + route_id: str + params: LiteLLMParamsBody + + +ANOMALY_ROUTES = ( + AnomalyRoute( + route_id="anthropic", + params=LiteLLMParamsBody(model="anthropic/claude-sonnet-5"), + ), + AnomalyRoute( + route_id="bedrock_invoke", + params=LiteLLMParamsBody( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + aws_region_name="us-east-1", + ), + ), +) + + +def _route_id(route: AnomalyRoute) -> str: + return route.route_id + + +def _settled_key_spend(proxy: ProxyClient, key: str) -> float: + deadline = time.monotonic() + proxy.poll_timeout + + def settle(previous: float) -> float: + current = proxy.key_info(key).spend or 0.0 + if current > 0 and current == previous: + return current + if time.monotonic() >= deadline: + raise AssertionError( + f"key spend never settled to a stable non-zero value within " + f"{proxy.poll_timeout}s (last read {current}); spend stopped being " + f"recorded, which is itself a spend anomaly" + ) + time.sleep(proxy.poll_interval) + return settle(current) + + return settle(-1.0) + + +class TestWeeklySessionAnomaly: + @pytest.mark.covers("reliability.perf.session_anomaly.under_slo") + @pytest.mark.parametrize("route", ANOMALY_ROUTES, ids=_route_id) + def test_session_load_stays_within_baselines( + self, client: LoadClient, resources: ResourceManager, route: AnomalyRoute + ) -> None: + model_name = f"weekly-anomaly-{route.route_id}-{unique_marker()}" + model_id = client.proxy.create_model(model_name, route.params) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = client.proxy.generate_key( + KeyGenerateBody(models=[model_name], key_alias=model_name) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + turns = run_concurrent_sessions( + client.proxy.transport, + key, + model_name, + ANOMALY_SESSIONS, + ANOMALY_TURNS_PER_SESSION, + ) + report = summarize(turns) + failures = tuple(turn.failure for turn in turns if turn.failure) + print(f"{route.route_id} anomaly report: {report}") + + assert report.error_ratio <= ANOMALY_MAX_ERROR_RATIO, ( + f"{route.route_id}: {report.failed_turns}/{report.attempted_turns} turns " + f"failed ({report.error_ratio:.1%} > {ANOMALY_MAX_ERROR_RATIO:.1%} allowed); " + f"error rate is anomalously high. Failures: {failures}" + ) + assert report.warm_turns > 0, ( + f"{route.route_id}: no session got past its first turn, so cache and " + f"latency baselines have nothing to read. Failures: {failures}" + ) + assert report.warm_cache_read_share >= ANOMALY_MIN_WARM_CACHE_READ_SHARE, ( + f"{route.route_id}: warm turns read only {report.warm_cache_read_share:.1%} " + f"of billed input tokens from the prompt cache " + f"(read={report.warm_cache_read_tokens}, " + f"creation={report.warm_cache_creation_tokens}, " + f"uncached={report.warm_uncached_input_tokens}), below the " + f"{ANOMALY_MIN_WARM_CACHE_READ_SHARE:.0%} floor; the cached prefix is " + f"being invalidated between turns (the mid-conversation-system cache " + f"collapse signature) or caching stopped working" + ) + assert report.warm_cache_creation_tokens > 0, ( + f"{route.route_id}: warm turns wrote 0 cache-creation tokens across " + f"{report.warm_turns} turns; the moving cache breakpoint stopped writing " + f"new prefix increments" + ) + assert report.p95_turn_seconds <= ANOMALY_MAX_P95_TURN_SECONDS, ( + f"{route.route_id}: p95 turn time {report.p95_turn_seconds:.1f}s exceeds " + f"the {ANOMALY_MAX_P95_TURN_SECONDS:.0f}s ceiling under " + f"{ANOMALY_SESSIONS} concurrent sessions; turn times are anomalously slow" + ) + + spend = _settled_key_spend(client.proxy, key) + assert spend <= ANOMALY_MAX_KEY_SPEND_USD, ( + f"{route.route_id}: gateway recorded ${spend:.4f} for " + f"{report.attempted_turns} turns, above the " + f"${ANOMALY_MAX_KEY_SPEND_USD} ceiling; spend per session is " + f"anomalously high (cache regressions surface here as 2-3x spend)" + ) diff --git a/tests/e2e/load/weekly_anomaly_config.yml b/tests/e2e/load/weekly_anomaly_config.yml new file mode 100644 index 00000000000..08972969cf0 --- /dev/null +++ b/tests/e2e/load/weekly_anomaly_config.yml @@ -0,0 +1,3 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index e9611df139b..2998a4b83c6 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -6,3 +6,4 @@ addopts = --strict-markers --strict-config markers = e2e: live test that requires a running proxy and real provider keys load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites + weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set From 51f0f40c2f34b35668730c4ab1f569213a3637a3 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 21 Jul 2026 22:33:07 +0000 Subject: [PATCH 008/130] test(e2e): invoke a real published a2a agent and assert it replies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/a2a/a2a_client.py | 88 ++++++++++++++++++++++++-- tests/e2e/a2a/test_a2a_agent_e2e.py | 72 ++++++++++----------- tests/e2e/coverage_registry/other.yaml | 2 +- 3 files changed, 115 insertions(+), 47 deletions(-) diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index c1a1503ec96..07bc2840b2b 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -11,6 +11,7 @@ here because only this suite uses them. from __future__ import annotations +import urllib.request import warnings from dataclasses import dataclass @@ -33,6 +34,11 @@ class A2ASkill(BaseModel): examples: list[str] | None = None +class A2AProvider(BaseModel): + organization: str + url: str + + class AgentCardParams(BaseModel): """The upstream agent card an admin registers. `protocolVersion` is the field the proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration.""" @@ -49,6 +55,29 @@ class AgentCardParams(BaseModel): preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport") +class UpstreamAgentCard(BaseModel): + """A real published agent card parsed from a public /.well-known endpoint. Keys on + the A2A wire aliases so `model_validate_json` reads the served JSON and + `model_dump(by_alias=True)` re-emits it unchanged for verbatim registration; it is + only ever fetched-and-validated, never hand-constructed, so aliasing on the wire + names does not affect any call site.""" + + model_config = ConfigDict(populate_by_name=True) + + protocol_version: str = Field(alias="protocolVersion") + name: str + description: str + version: str + url: str + provider: A2AProvider | None = None + documentation_url: str | None = Field(default=None, alias="documentationUrl") + capabilities: A2ACapabilities = A2ACapabilities() + skills: list[A2ASkill] + default_input_modes: list[str] = Field(default=["text"], alias="defaultInputModes") + default_output_modes: list[str] = Field(default=["text"], alias="defaultOutputModes") + preferred_transport: str | None = Field(default=None, alias="preferredTransport") + + class A2ABridgeParams(BaseModel): """litellm_params that route the agent through the completion bridge: an A2A message/send is transformed into a litellm.acompletion against this provider.""" @@ -61,8 +90,8 @@ class A2ABridgeParams(BaseModel): class AgentRegisterBody(BaseModel): agent_name: str - agent_card_params: AgentCardParams - litellm_params: A2ABridgeParams + agent_card_params: AgentCardParams | UpstreamAgentCard + litellm_params: A2ABridgeParams | None = None class A2ASecurityScheme(BaseModel): @@ -104,9 +133,32 @@ class A2ATextPart(BaseModel): text: str +class A2ASearchPropertiesParams(BaseModel): + """The strict param schema of the published property agent's `search_properties` + skill (unknown keys are rejected upstream), so a natural-language query like + "properties for sale in SF under $2M" is expressed as typed fields.""" + + un_locode: str | None = None + service_type: str | None = None + property_type: str | None = None + bedrooms_min: int | None = None + asking_price_max: float | None = None + limit: int | None = None + + +class A2ASkillInvocation(BaseModel): + skill: str + params: A2ASearchPropertiesParams + + +class A2ADataPart(BaseModel): + kind: str = "data" + data: A2ASkillInvocation + + class A2AOutboundMessage(BaseModel): role: str = "user" - parts: list[A2ATextPart] + parts: list[A2ATextPart | A2ADataPart] message_id: str = Field(serialization_alias="messageId") @@ -134,10 +186,16 @@ class A2AResponseMessage(BaseModel): parts: list[A2AResponsePart] = [] +class A2ATaskStatus(BaseModel): + state: str | None = None + message: A2AResponseMessage | None = None + + class A2AResult(BaseModel): """A message/send result. In 0.3 the message fields sit directly on the result - (`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the - agent's reply from whichever shape the served version produced.""" + (`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that + runs a task replies with a `task` whose agent text lives on `status.message`. + `text` reads the agent's reply from whichever shape the served version produced.""" model_config = ConfigDict(populate_by_name=True) @@ -146,10 +204,18 @@ class A2AResult(BaseModel): message_id: str | None = Field(default=None, alias="messageId") parts: list[A2AResponsePart] = [] message: A2AResponseMessage | None = None + status: A2ATaskStatus | None = None @property def text(self) -> str: - parts = self.message.parts if self.message is not None else self.parts + if self.message is not None: + parts = self.message.parts + elif self.parts: + parts = self.parts + elif self.status is not None and self.status.message is not None: + parts = self.status.message.parts + else: + parts = [] return "".join(part.text or "" for part in parts) @property @@ -218,3 +284,13 @@ class A2AClient: def build_a2a_client(proxy: ProxyClient) -> A2AClient: return A2AClient(proxy=proxy) + + +def fetch_agent_card(url: str, *, timeout: float = 20.0) -> UpstreamAgentCard: + """Fetch a live A2A agent card from its /.well-known endpoint and parse it into the + registration model, so a test can register a real published card verbatim rather + than a hand-rolled one.""" + request = urllib.request.Request(url, headers={"Accept": "application/json"}) + with urllib.request.urlopen(request, timeout=timeout) as response: # noqa: S310 # pyright: ignore[reportAny] # fixed https well-known url; typeshed types urlopen as Any + payload: bytes = response.read() # pyright: ignore[reportAny] # typeshed types urlopen as Any + return UpstreamAgentCard.model_validate_json(payload) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index 8c19f97e64c..1db7f63ca54 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -14,16 +14,19 @@ import pytest from a2a_client import ( A2ABridgeParams, - A2ACapabilities, A2AClient, + A2ADataPart, A2AJsonRpcRequest, A2AMessageSendParams, A2AOutboundMessage, + A2ASearchPropertiesParams, A2ASkill, + A2ASkillInvocation, A2ATextPart, AgentCardParams, AgentRegisterBody, AgentResponse, + fetch_agent_card, ) from e2e_config import unique_marker from e2e_http import Result, UnknownApiError, unwrap @@ -31,6 +34,9 @@ from lifecycle import ResourceManager BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5") +MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json" +MOVEHOME_ORIGIN = "https://movehome.org" + pytestmark = pytest.mark.e2e @@ -52,34 +58,6 @@ def _register(client: A2AClient, resources: ResourceManager, protocol_version: s return agent -def _google_sdk_default_card(marker: str) -> AgentCardParams: - """Field-for-field the card the Google a2a-sdk 0.3.x line serves for its helloworld - sample, whose ``AgentCard`` defaults ``protocolVersion`` to the full semver "0.3.0" - (the shape the v1.92 regression rejected); ``url`` and ``name`` are the - deployment-specific fields the SDK requires callers to fill. The url points at a - reserved example host: a url-bearing card makes message/send dial that upstream - rather than the completion bridge, so the verbatim-card test asserts registration - and card serving only.""" - return AgentCardParams( - protocol_version="0.3.0", - name=f"E2E A2A sdk {marker}", - description="Just a hello world agent", - version="1.0.0", - url="http://e2e-a2a-upstream.example/", - capabilities=A2ACapabilities(streaming=True), - skills=[ - A2ASkill( - id="hello_world", - name="Returns hello world", - description="just returns hello world", - tags=["hello world"], - examples=["hi", "hello world"], - ) - ], - preferred_transport="JSONRPC", - ) - - def _register_rejection(client: A2AClient, protocol_version: str) -> Result[AgentResponse]: marker = unique_marker() body = AgentRegisterBody( @@ -126,21 +104,35 @@ class TestA2AAgentLifecycle: assert result is not None assert result.text != "" - @pytest.mark.covers("other.a2a.register.sdk_default_card_accepted") - def test_google_sdk_default_card_registers_verbatim(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + @pytest.mark.covers("other.a2a.message_send.real_world_agent_replies") + def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + upstream = fetch_agent_card(MOVEHOME_AGENT_CARD_URL).model_copy(update={"url": MOVEHOME_ORIGIN}) + assert upstream.protocol_version == "0.3.0" marker = unique_marker() - body = AgentRegisterBody( - agent_name=f"e2e-a2a-sdk-{marker}", - agent_card_params=_google_sdk_default_card(marker), - litellm_params=BRIDGE, - ) + body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream) agent = unwrap(client.register_agent(body)) resources.defer(lambda: client.delete_agent(agent.agent_id)) assert agent.agent_card_params.protocol_version == "0.3" - card = unwrap(client.agent_card(agent.agent_id, scoped_key)) - assert card.protocol_version == "0.3" - assert card.supported_interfaces is not None - assert card.supported_interfaces[0].protocol_version == "0.3" + request = A2AJsonRpcRequest( + id=f"e2e-{unique_marker()}", + params=A2AMessageSendParams( + message=A2AOutboundMessage( + parts=[ + A2ADataPart( + data=A2ASkillInvocation( + skill="search_properties", + params=A2ASearchPropertiesParams(un_locode="USSFO", service_type="sale", asking_price_max=2_000_000, limit=3), + ) + ) + ], + message_id=unique_marker(), + ) + ), + ) + response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) + assert response.error is None + assert response.result is not None + assert response.result.text.strip() != "" @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 00a76cf511a..ace4f8bcdc9 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -31,7 +31,7 @@ - {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} - {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} - {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"} -- {id: other.a2a.register.sdk_default_card_accepted, module: other, tier: P1, area: a2a, assertions: [sdk_default_card_accepted], source: "agent_endpoints/endpoints.py _build_merged_agent_card", rationale: "The full field set the Google a2a-sdk 0.3.x emits (protocolVersion 0.3.0, url, preferredTransport, capabilities, skill examples) registers verbatim and serves the canonical 0.3"} +- {id: other.a2a.message_send.real_world_agent_replies, module: other, tier: P1, area: a2a, assertions: [real_world_agent_replies], source: "agent_endpoints/a2a_endpoints.py asend_message", rationale: "A real published a2a agent fetched live from a public /.well-known endpoint (pinning the full semver 0.3.0 the a2a-sdk emits) registers, serves the canonical 0.3, and a message/send skill invocation proxies to the live upstream and returns the agent's reply"} - {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"} - {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} From 1692170264102008ab8fd7e1686af8adc481331c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:34:56 -0700 Subject: [PATCH 009/130] fix(e2e): count aborted-session turns as failures and require a spend stability window --- tests/e2e/e2e_config.py | 3 + tests/e2e/load/session_anomaly.py | 37 ++++++- tests/e2e/load/test_session_anomaly.py | 98 +++++++++++++++++++ .../load/test_weekly_session_anomaly_e2e.py | 33 +++---- 4 files changed, 147 insertions(+), 24 deletions(-) create mode 100644 tests/e2e/load/test_session_anomaly.py diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 985b299eac8..1b75c8e8ce1 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -91,6 +91,9 @@ ANOMALY_MAX_P95_TURN_SECONDS = float( ANOMALY_MAX_KEY_SPEND_USD = float( os.environ.get("E2E_ANOMALY_MAX_KEY_SPEND_USD", "0.60") ) +ANOMALY_SPEND_SETTLE_SECONDS = float( + os.environ.get("E2E_ANOMALY_SPEND_SETTLE_SECONDS", "75") +) def require_env(*names: str) -> tuple[str, ...]: diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/load/session_anomaly.py index fc1cfa7ac70..b2e23d24ab2 100644 --- a/tests/e2e/load/session_anomaly.py +++ b/tests/e2e/load/session_anomaly.py @@ -1,6 +1,7 @@ from __future__ import annotations import time +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass @@ -53,6 +54,7 @@ class TurnMetric: @dataclass(frozen=True, slots=True) class AnomalyReport: + planned_turns: int attempted_turns: int failed_turns: int warm_turns: int @@ -63,7 +65,7 @@ class AnomalyReport: @property def error_ratio(self) -> float: - return self.failed_turns / self.attempted_turns if self.attempted_turns else 1.0 + return self.failed_turns / self.planned_turns if self.planned_turns else 1.0 @property def warm_cache_read_share(self) -> float: @@ -217,6 +219,34 @@ def run_concurrent_sessions( return tuple(turn for future in futures for turn in future.result()) +def settled_spend( + read_spend: Callable[[], float], + poll_interval: float, + settle_seconds: float, + timeout_seconds: float, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> float: + deadline = now() + timeout_seconds + settle_seconds + + def settle(previous: float, stable_since: float) -> float: + current = read_spend() + observed = now() + since = stable_since if current == previous else observed + if current > 0 and observed - since >= settle_seconds: + return current + if observed >= deadline: + raise AssertionError( + f"key spend never held a stable non-zero value for {settle_seconds}s " + f"within {timeout_seconds + settle_seconds}s (last read {current}); " + f"spend stopped being recorded, which is itself a spend anomaly" + ) + sleep(poll_interval) + return settle(current, since) + + return settle(-1.0, now()) + + def _p95(latencies: tuple[float, ...]) -> float: if not latencies: return 0.0 @@ -224,11 +254,12 @@ def _p95(latencies: tuple[float, ...]) -> float: return ranked[max(0, -(-len(ranked) * 95 // 100) - 1)] -def summarize(turns: tuple[TurnMetric, ...]) -> AnomalyReport: +def summarize(turns: tuple[TurnMetric, ...], planned_turns: int) -> AnomalyReport: warm = tuple(turn for turn in turns if turn.ok and turn.turn_index >= 2) return AnomalyReport( + planned_turns=planned_turns, attempted_turns=len(turns), - failed_turns=sum(1 for turn in turns if not turn.ok), + failed_turns=planned_turns - sum(1 for turn in turns if turn.ok), warm_turns=len(warm), warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm), warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm), diff --git a/tests/e2e/load/test_session_anomaly.py b/tests/e2e/load/test_session_anomaly.py new file mode 100644 index 00000000000..9acedbac424 --- /dev/null +++ b/tests/e2e/load/test_session_anomaly.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from itertools import count, repeat + +import pytest + +from session_anomaly import TurnMetric, settled_spend, summarize + + +def _ok_turn(turn_index: int) -> TurnMetric: + return TurnMetric( + turn_index=turn_index, + ok=True, + latency_seconds=1.0, + uncached_input_tokens=10, + cache_read_tokens=100, + cache_creation_tokens=5, + failure=None, + ) + + +def _failed_turn(turn_index: int) -> TurnMetric: + return TurnMetric( + turn_index=turn_index, + ok=False, + latency_seconds=1.0, + uncached_input_tokens=0, + cache_read_tokens=0, + cache_creation_tokens=0, + failure="NetworkError()", + ) + + +class TestSummarizePlannedTurns: + def test_session_aborted_on_first_turn_counts_all_its_planned_turns_as_failed( + self, + ) -> None: + completed_session = tuple(_ok_turn(index) for index in range(1, 7)) + aborted_session = (_failed_turn(1),) + + report = summarize((*completed_session, *aborted_session), planned_turns=12) + + assert report.attempted_turns == 7 + assert report.failed_turns == 6 + assert report.error_ratio == 0.5 + + def test_all_planned_turns_completing_reports_zero_failures(self) -> None: + report = summarize( + tuple(_ok_turn(index) for index in range(1, 7)), planned_turns=6 + ) + + assert report.failed_turns == 0 + assert report.error_ratio == 0.0 + + +class TestSettledSpend: + def test_partial_total_between_batch_flushes_is_not_accepted_as_final(self) -> None: + reads = iter((0.1, 0.1, 0.1, 0.35, 0.35, 0.35, 0.35, 0.35)) + ticks = count(0.0, 2.5) + + spend = settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=10.0, + timeout_seconds=100.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) + + assert spend == 0.35 + + def test_spend_that_never_stabilizes_raises(self) -> None: + reads = (0.1 * step for step in count(1)) + ticks = count(0.0, 2.5) + + with pytest.raises(AssertionError, match="spend anomaly"): + settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=5.0, + timeout_seconds=10.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) + + def test_spend_that_never_becomes_nonzero_raises(self) -> None: + reads = repeat(0.0) + ticks = count(0.0, 2.5) + + with pytest.raises(AssertionError, match="spend anomaly"): + settled_spend( + lambda: next(reads), + poll_interval=5.0, + settle_seconds=5.0, + timeout_seconds=10.0, + now=lambda: next(ticks), + sleep=lambda _: None, + ) diff --git a/tests/e2e/load/test_weekly_session_anomaly_e2e.py b/tests/e2e/load/test_weekly_session_anomaly_e2e.py index ebd58166f0b..de7ca33ba7c 100644 --- a/tests/e2e/load/test_weekly_session_anomaly_e2e.py +++ b/tests/e2e/load/test_weekly_session_anomaly_e2e.py @@ -1,6 +1,5 @@ from __future__ import annotations -import time from dataclasses import dataclass import pytest @@ -11,6 +10,7 @@ from e2e_config import ( ANOMALY_MAX_P95_TURN_SECONDS, ANOMALY_MIN_WARM_CACHE_READ_SHARE, ANOMALY_SESSIONS, + ANOMALY_SPEND_SETTLE_SECONDS, ANOMALY_TURNS_PER_SESSION, unique_marker, ) @@ -18,7 +18,7 @@ from lifecycle import ResourceManager from load_client import LoadClient from models import KeyGenerateBody, LiteLLMParamsBody from proxy_client import ProxyClient -from session_anomaly import run_concurrent_sessions, summarize +from session_anomaly import run_concurrent_sessions, settled_spend, summarize pytestmark = [pytest.mark.e2e, pytest.mark.load, pytest.mark.weekly] @@ -49,22 +49,12 @@ def _route_id(route: AnomalyRoute) -> str: def _settled_key_spend(proxy: ProxyClient, key: str) -> float: - deadline = time.monotonic() + proxy.poll_timeout - - def settle(previous: float) -> float: - current = proxy.key_info(key).spend or 0.0 - if current > 0 and current == previous: - return current - if time.monotonic() >= deadline: - raise AssertionError( - f"key spend never settled to a stable non-zero value within " - f"{proxy.poll_timeout}s (last read {current}); spend stopped being " - f"recorded, which is itself a spend anomaly" - ) - time.sleep(proxy.poll_interval) - return settle(current) - - return settle(-1.0) + return settled_spend( + lambda: proxy.key_info(key).spend or 0.0, + proxy.poll_interval, + ANOMALY_SPEND_SETTLE_SECONDS, + proxy.poll_timeout, + ) class TestWeeklySessionAnomaly: @@ -88,13 +78,14 @@ class TestWeeklySessionAnomaly: ANOMALY_SESSIONS, ANOMALY_TURNS_PER_SESSION, ) - report = summarize(turns) + report = summarize(turns, ANOMALY_SESSIONS * ANOMALY_TURNS_PER_SESSION) failures = tuple(turn.failure for turn in turns if turn.failure) print(f"{route.route_id} anomaly report: {report}") assert report.error_ratio <= ANOMALY_MAX_ERROR_RATIO, ( - f"{route.route_id}: {report.failed_turns}/{report.attempted_turns} turns " - f"failed ({report.error_ratio:.1%} > {ANOMALY_MAX_ERROR_RATIO:.1%} allowed); " + f"{route.route_id}: {report.failed_turns}/{report.planned_turns} planned " + f"turns failed or never ran because their session aborted " + f"({report.error_ratio:.1%} > {ANOMALY_MAX_ERROR_RATIO:.1%} allowed); " f"error rate is anomalously high. Failures: {failures}" ) assert report.warm_turns > 0, ( From 48295df0e93af82ba243ed7792e2e8c721fbf3e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:21:45 -0700 Subject: [PATCH 010/130] test(e2e): raise warm cache read floor to 0.65 from measured healthy and regression baselines --- tests/e2e/e2e_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 1b75c8e8ce1..4fdd91aa3d9 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -83,7 +83,7 @@ ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_MAX_ERROR_RATIO = float(os.environ.get("E2E_ANOMALY_MAX_ERROR_RATIO", "0.05")) ANOMALY_MIN_WARM_CACHE_READ_SHARE = float( - os.environ.get("E2E_ANOMALY_MIN_WARM_CACHE_READ_SHARE", "0.5") + os.environ.get("E2E_ANOMALY_MIN_WARM_CACHE_READ_SHARE", "0.65") ) ANOMALY_MAX_P95_TURN_SECONDS = float( os.environ.get("E2E_ANOMALY_MAX_P95_TURN_SECONDS", "30") From 0bfdb3726675234d49362702a8d21dc99d576db4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:35:02 -0700 Subject: [PATCH 011/130] test(e2e): route external agent card fetch through the typed transport Adds get_external to e2e_http.py for absolute third-party GETs (no proxy base url or auth, same Result classification) and rewires fetch_agent_card through it, dropping the urllib.request escape hatch. Creates tests/code_coverage_tests/check_e2e_no_raw_requests.py, the checker tests/e2e/CLAUDE.md already referenced, and wires it into the code-quality workflow so raw HTTP client imports outside the transport fail CI; pre-existing uses (root conftest liveness probe, claude_code version resolver) are grandfathered and exception-type-only imports stay allowed. --- .github/workflows/test-code-quality.yml | 3 + .../check_e2e_no_raw_requests.py | 81 +++++++++++++++++++ tests/e2e/a2a/a2a_client.py | 10 +-- tests/e2e/a2a/test_a2a_agent_e2e.py | 2 +- tests/e2e/e2e_http.py | 20 +++++ 5 files changed, 108 insertions(+), 8 deletions(-) create mode 100644 tests/code_coverage_tests/check_e2e_no_raw_requests.py diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 9d28ca211cf..ae31395521a 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -115,6 +115,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_e2e_no_raw_requests + run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py + - name: memory_test run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py new file mode 100644 index 00000000000..e70e83652d1 --- /dev/null +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -0,0 +1,81 @@ +"""tests/e2e routes every HTTP call through the typed transport (e2e_http.py), so +raw HTTP client imports (requests, urllib.request, httpx, aiohttp, http.client) are +banned in suite code. Importing requests' exception types for catching is fine +anywhere; a small allowlist grandfathers the files that legitimately make raw calls +(the transport itself, the root conftest liveness probe, and the claude_code version +resolver's constant registry URL fetch). Referenced by tests/e2e/CLAUDE.md.""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +E2E_DIR = Path(__file__).resolve().parents[1] / "e2e" + +BANNED_MODULES = ("requests", "urllib.request", "http.client", "httpx", "aiohttp") + +ALLOWED_RAW_CLIENT_FILES = { + "e2e_http.py": ("requests",), + "conftest.py": ("requests",), + "claude_code/pr_gate_version_resolver.py": ("urllib.request",), +} + +EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) + + +def _is_banned(module: str) -> bool: + return any(module == banned or module.startswith(banned + ".") for banned in BANNED_MODULES) + + +def _banned_imports(tree: ast.Module) -> tuple[tuple[str, int], ...]: + plain = tuple( + (alias.name, node.lineno) + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + if _is_banned(alias.name) + ) + from_imports = tuple( + (node.module, node.lineno) + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) + and node.module is not None + and _is_banned(node.module) + and not all(alias.name in EXCEPTION_ONLY_NAMES for alias in node.names) + ) + return plain + from_imports + + +def _violations_in(path: Path) -> tuple[str, ...]: + relative = path.relative_to(E2E_DIR).as_posix() + allowed = ALLOWED_RAW_CLIENT_FILES.get(relative, ()) + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + f"tests/e2e/{relative}:{lineno}: raw HTTP client import '{module}'" + for module, lineno in _banned_imports(tree) + if module not in allowed + ) + + +def main() -> int: + violations = tuple( + violation + for path in sorted(E2E_DIR.rglob("*.py")) + for violation in _violations_in(path) + ) + for violation in violations: + print(violation) + if violations: + print( + f"\n{len(violations)} raw HTTP client import(s) in tests/e2e. " + "Route the call through tests/e2e/e2e_http.py (get_external for absolute " + "third-party URLs) so it gets the typed Result handling." + ) + return 1 + print("tests/e2e raw HTTP client check passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index 07bc2840b2b..97ffa8c34a3 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -11,13 +11,12 @@ here because only this suite uses them. from __future__ import annotations -import urllib.request import warnings from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field -from e2e_http import NoBody, Result, is_ok +from e2e_http import NoBody, Result, get_external, is_ok from proxy_client import ProxyClient @@ -286,11 +285,8 @@ def build_a2a_client(proxy: ProxyClient) -> A2AClient: return A2AClient(proxy=proxy) -def fetch_agent_card(url: str, *, timeout: float = 20.0) -> UpstreamAgentCard: +def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]: """Fetch a live A2A agent card from its /.well-known endpoint and parse it into the registration model, so a test can register a real published card verbatim rather than a hand-rolled one.""" - request = urllib.request.Request(url, headers={"Accept": "application/json"}) - with urllib.request.urlopen(request, timeout=timeout) as response: # noqa: S310 # pyright: ignore[reportAny] # fixed https well-known url; typeshed types urlopen as Any - payload: bytes = response.read() # pyright: ignore[reportAny] # typeshed types urlopen as Any - return UpstreamAgentCard.model_validate_json(payload) + return get_external(url, response_type=UpstreamAgentCard, timeout=timeout) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index 1db7f63ca54..aa60b57f99b 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -106,7 +106,7 @@ class TestA2AAgentLifecycle: @pytest.mark.covers("other.a2a.message_send.real_world_agent_replies") def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: - upstream = fetch_agent_card(MOVEHOME_AGENT_CARD_URL).model_copy(update={"url": MOVEHOME_ORIGIN}) + upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN}) assert upstream.protocol_version == "0.3.0" marker = unique_marker() body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 1c6048a3688..96c1fb6a6b7 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -244,6 +244,26 @@ def get[R: BaseModel]( return _classify(resp, response_type) +def get_external[R: BaseModel]( + url: str, + *, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + """GET an absolute URL outside the proxy (e.g. a public /.well-known document). + Unlike the transport wrappers there is no proxy base url and no proxy auth; the + response still gets the same tagged-union classification as every other call.""" + try: + resp = requests.get( + url, + headers={"Accept": "application/json"}, + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + def delete[R: BaseModel]( url: URL, *, From 560dc6dd0d22b887ede19371e03f1e5fe9c8d45d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:47:47 -0700 Subject: [PATCH 012/130] ci: run check_e2e_no_raw_requests in make pre-commit for staged tests/e2e files Mirrors the new test-code-quality.yml step locally so a green pre-commit stays predictive: the sub-second checker fires only when tests/e2e Python files are staged, matching the script's staged-file gating for every other block. --- scripts/pre_commit_lint.sh | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index cce0cb61c1e..150a4bbf9de 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -5,6 +5,7 @@ # gating CI checks, so a clean run means a green CI lint: # - litellm/ Python staged -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) +# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) # - dashboard staged -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) # @@ -112,6 +113,12 @@ if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; } fi +if [ -n "$e2e_py_files" ]; then + echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" + uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \ + || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; } +fi + if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then From 1255094de30a0bcaf1dc0c9a46e86337ed0191cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:55:44 -0700 Subject: [PATCH 013/130] fix(e2e): retry transient turn failures and move the weekly anomaly run to Saturday before the stable release cut --- .github/workflows/weekly_load_anomaly.yml | 2 +- tests/e2e/e2e_config.py | 1 + tests/e2e/load/session_anomaly.py | 51 +++++++++++++++---- tests/e2e/load/test_session_anomaly.py | 45 +++++++++++++++- .../load/test_weekly_session_anomaly_e2e.py | 2 + 5 files changed, 88 insertions(+), 13 deletions(-) diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 6853ffa1cde..4c2103f026d 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -2,7 +2,7 @@ name: "Weekly Load Anomaly Check" on: schedule: - - cron: "0 6 * * 1" + - cron: "0 12 * * 6" workflow_dispatch: permissions: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 4fdd91aa3d9..30353b93dd6 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -81,6 +81,7 @@ LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.0 WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) +ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) ANOMALY_MAX_ERROR_RATIO = float(os.environ.get("E2E_ANOMALY_MAX_ERROR_RATIO", "0.05")) ANOMALY_MIN_WARM_CACHE_READ_SHARE = float( os.environ.get("E2E_ANOMALY_MIN_WARM_CACHE_READ_SHARE", "0.65") diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/load/session_anomaly.py index b2e23d24ab2..c29b833635d 100644 --- a/tests/e2e/load/session_anomaly.py +++ b/tests/e2e/load/session_anomaly.py @@ -110,6 +110,22 @@ def _without_cache_control(message: RichMessage) -> RichMessage: ) +RETRY_BACKOFF_SECONDS = 2.0 + + +def retried( + call: Callable[[], Result[SessionMessagesResponse]], + attempts: int, + backoff_seconds: float = RETRY_BACKOFF_SECONDS, + sleep: Callable[[float], None] = time.sleep, +) -> Result[SessionMessagesResponse]: + result = call() + if isinstance(result, Success) or attempts <= 1: + return result + sleep(backoff_seconds) + return retried(call, attempts - 1, backoff_seconds, sleep) + + def _metric( result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float ) -> TurnMetric: @@ -144,6 +160,7 @@ def _drive_turns( history: tuple[RichMessage, ...], turn_index: int, remaining_turns: int, + attempts_per_turn: int, ) -> tuple[TurnMetric, ...]: if remaining_turns == 0: return () @@ -156,15 +173,18 @@ def _drive_turns( ], ) started = time.monotonic() - result = transport.post( - "/v1/messages", - headers=transport.bearer(key), - json=SessionMessagesRequest( - model=model, - system=[system_block], - messages=[*history, user_turn], + result = retried( + lambda: transport.post( + "/v1/messages", + headers=transport.bearer(key), + json=SessionMessagesRequest( + model=model, + system=[system_block], + messages=[*history, user_turn], + ), + response_type=SessionMessagesResponse, ), - response_type=SessionMessagesResponse, + attempts_per_turn, ) turn = _metric(result, turn_index, time.monotonic() - started) if not isinstance(result, Success): @@ -188,12 +208,13 @@ def _drive_turns( ), turn_index + 1, remaining_turns - 1, + attempts_per_turn, ), ) def run_session( - transport: Transport, key: str, model: str, turns: int + transport: Transport, key: str, model: str, turns: int, attempts_per_turn: int ) -> tuple[TurnMetric, ...]: marker = unique_marker() return _drive_turns( @@ -205,15 +226,23 @@ def run_session( (), 1, turns, + attempts_per_turn, ) def run_concurrent_sessions( - transport: Transport, key: str, model: str, sessions: int, turns_per_session: int + transport: Transport, + key: str, + model: str, + sessions: int, + turns_per_session: int, + attempts_per_turn: int, ) -> tuple[TurnMetric, ...]: with ThreadPoolExecutor(max_workers=sessions) as pool: futures = [ - pool.submit(run_session, transport, key, model, turns_per_session) + pool.submit( + run_session, transport, key, model, turns_per_session, attempts_per_turn + ) for _ in range(sessions) ] return tuple(turn for future in futures for turn in future.result()) diff --git a/tests/e2e/load/test_session_anomaly.py b/tests/e2e/load/test_session_anomaly.py index 9acedbac424..80f8aff3ba4 100644 --- a/tests/e2e/load/test_session_anomaly.py +++ b/tests/e2e/load/test_session_anomaly.py @@ -4,7 +4,14 @@ from itertools import count, repeat import pytest -from session_anomaly import TurnMetric, settled_spend, summarize +from e2e_http import NetworkError, Success +from session_anomaly import ( + SessionMessagesResponse, + TurnMetric, + retried, + settled_spend, + summarize, +) def _ok_turn(turn_index: int) -> TurnMetric: @@ -53,6 +60,42 @@ class TestSummarizePlannedTurns: assert report.error_ratio == 0.0 +class TestRetried: + def test_transient_failures_then_success_returns_the_success(self) -> None: + outcome = Success(data=SessionMessagesResponse()) + calls = iter( + (NetworkError(message="overloaded"), NetworkError(message="overloaded"), outcome) + ) + + result = retried(lambda: next(calls), attempts=3, sleep=lambda _: None) + + assert result is outcome + + def test_exhausted_attempts_return_the_last_failure(self) -> None: + last_attempt = NetworkError(message="still overloaded") + never_reached = NetworkError(message="a fourth attempt would break the budget") + calls = iter( + (NetworkError(message="overloaded"), last_attempt, never_reached) + ) + + result = retried(lambda: next(calls), attempts=2, sleep=lambda _: None) + + assert result is last_attempt + assert next(calls) is never_reached + + def test_first_try_success_never_sleeps(self) -> None: + def sleep_means_retry(_: float) -> None: + raise AssertionError("slept after a successful attempt") + + result = retried( + lambda: Success(data=SessionMessagesResponse()), + attempts=3, + sleep=sleep_means_retry, + ) + + assert isinstance(result, Success) + + class TestSettledSpend: def test_partial_total_between_batch_flushes_is_not_accepted_as_final(self) -> None: reads = iter((0.1, 0.1, 0.1, 0.35, 0.35, 0.35, 0.35, 0.35)) diff --git a/tests/e2e/load/test_weekly_session_anomaly_e2e.py b/tests/e2e/load/test_weekly_session_anomaly_e2e.py index de7ca33ba7c..d4ef883702e 100644 --- a/tests/e2e/load/test_weekly_session_anomaly_e2e.py +++ b/tests/e2e/load/test_weekly_session_anomaly_e2e.py @@ -11,6 +11,7 @@ from e2e_config import ( ANOMALY_MIN_WARM_CACHE_READ_SHARE, ANOMALY_SESSIONS, ANOMALY_SPEND_SETTLE_SECONDS, + ANOMALY_TURN_ATTEMPTS, ANOMALY_TURNS_PER_SESSION, unique_marker, ) @@ -77,6 +78,7 @@ class TestWeeklySessionAnomaly: model_name, ANOMALY_SESSIONS, ANOMALY_TURNS_PER_SESSION, + ANOMALY_TURN_ATTEMPTS, ) report = summarize(turns, ANOMALY_SESSIONS * ANOMALY_TURNS_PER_SESSION) failures = tuple(turn.failure for turn in turns if turn.failure) From e4343eb14819cbbb1ba1b049ef53bb0db55693fa Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 21 Jul 2026 17:36:33 -0700 Subject: [PATCH 014/130] feat(rust): honor pre-computed Entra ID auth for Azure /messages (#34107) * feat(rust): honor pre-computed Entra ID (Authorization: Bearer) auth for Azure /messages * harden Rust Azure auth gate to require a non-empty Bearer token, not header presence --- .../ai-gateway/src/messages/common_utils.rs | 12 ++ .../crates/ai-gateway/src/messages/prepare.rs | 6 +- .../crates/ai-gateway/src/messages/tests.rs | 136 +++++++++++++++++- .../core/src/messages/transformation.rs | 4 + .../azure_ai/messages/transformation.rs | 9 ++ 5 files changed, 164 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index 4b906155665..68ecc3f17c1 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -50,3 +50,15 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { .iter() .any(|(key, _)| key.eq_ignore_ascii_case(name)) } + +pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + if !name.eq_ignore_ascii_case("authorization") { + return false; + } + let value = value.trim(); + value.len() > 7 + && value[..7].eq_ignore_ascii_case("bearer ") + && !value[7..].trim().is_empty() + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 624c3598fb0..9a027490eb6 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -3,7 +3,7 @@ use litellm_core::CoreResult; use litellm_core::messages::transformation::MessagesAuthStrategy; use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; -use super::common_utils::{has_header, messages_provider_config, string_headers}; +use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; use super::types::{MessagesRequest, ProviderMessagesRequest}; pub(super) fn prepare_messages_call( @@ -33,7 +33,9 @@ pub(super) fn prepare_messages_call( let mut headers = string_headers(request.extra_headers)?; let auth_strategy = config.auth_strategy(); - if !has_header(&headers, auth_strategy.header_name()) { + let already_authorized = has_header(&headers, auth_strategy.header_name()) + || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); + if !already_authorized { let api_key = config.resolve_api_key(request.api_key, &env_lookup)?; let auth_header = match auth_strategy { MessagesAuthStrategy::Bearer => { diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index a2d0f6fae23..23a53e98045 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -6,7 +6,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{ - has_header, messages_provider_config, string_headers, truncate_error_body, + has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, }; use super::{MessagesRequest, messages}; @@ -85,6 +85,34 @@ fn has_header_is_case_insensitive() { assert!(!has_header(&headers, "authorization")); } +#[test] +fn has_bearer_auth_requires_a_nonempty_bearer_token() { + assert!(has_bearer_auth(&[( + "Authorization".to_string(), + "Bearer tok".to_string() + )])); + assert!(has_bearer_auth(&[( + "authorization".to_string(), + "bearer tok".to_string() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + "Bearer ".to_string() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + String::new() + )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + "Basic abc".to_string() + )])); + assert!(!has_bearer_auth(&[( + "x-api-key".to_string(), + "sk".to_string() + )])); +} + #[tokio::test] async fn messages_round_trip_builds_azure_request_and_passes_response_through() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); @@ -252,6 +280,112 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { assert!(!head.contains("rust-fallback-key"), "{head}"); } +#[tokio::test] +async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = + r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer entra-token".to_string()), + ); + + messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: None, + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(headers), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("entra id request succeeds without api key"); + + let request = server.await.expect("server task completes"); + let head = request + .split_once("\r\n\r\n") + .expect("has body") + .0 + .to_ascii_lowercase(); + assert!(head.contains("authorization: bearer entra-token"), "{head}"); + assert!(!head.contains("x-api-key"), "{head}"); +} + +#[tokio::test] +async fn messages_requires_auth_when_no_key_and_no_header() { + let err = messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: None, + api_base: Some("http://127.0.0.1:1"), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + timeout: Some(Duration::from_millis(50)), + }) + .await + .expect_err("missing auth errors"); + + assert!(matches!(err, CoreError::Auth(_))); +} + +#[tokio::test] +async fn messages_ignores_malformed_authorization_and_uses_api_key() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = + r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer ".to_string()), + ); + + messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: Some("sk-azure"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(headers), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("falls back to api key"); + + let request = server.await.expect("server task completes"); + let head = request + .split_once("\r\n\r\n") + .expect("has body") + .0 + .to_ascii_lowercase(); + assert!(head.contains("x-api-key: sk-azure"), "{head}"); +} + #[tokio::test] async fn messages_maps_provider_error_status_to_http_error() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 3a34a58de6f..b478e20d24b 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -35,6 +35,10 @@ pub trait AnthropicMessagesProviderConfig: Sync { MessagesAuthStrategy::Header("x-api-key") } + fn accepts_bearer_auth(&self) -> bool { + false + } + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[ ("anthropic-version", "2023-06-01"), diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 6935bb4604b..7b958c77ba3 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -163,6 +163,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { self.anthropic.auth_strategy() } + fn accepts_bearer_auth(&self) -> bool { + true + } + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { self.anthropic.default_headers() } @@ -294,6 +298,11 @@ mod tests { ); } + #[test] + fn accepts_bearer_auth_for_entra_id() { + assert!(AZURE_ANTHROPIC_MESSAGES_CONFIG.accepts_bearer_auth()); + } + #[test] fn default_headers_match_python() { assert_eq!( From 1ea7db2111aa54ee50e03caf535e1237e01128a7 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 21 Jul 2026 17:36:47 -0700 Subject: [PATCH 015/130] fix(rust): route agentic-completion-hook /messages requests to Python for all stream modes (#34126) --- litellm/llms/custom_httpx/llm_http_handler.py | 8 +++----- .../anthropic_interface/test_rust_bridge_messages.py | 10 ++++------ 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c48d75439a7..ec1301e5923 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2107,8 +2107,7 @@ class BaseLLMHTTPHandler: rust_messages_response = await self._maybe_rust_anthropic_messages( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, - stream=stream or False, - rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj), + has_agentic_hook=self._has_agentic_completion_hook(logging_obj), model=model, api_key=api_key, api_base=api_base, @@ -2266,8 +2265,7 @@ class BaseLLMHTTPHandler: *, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, - stream: bool, - rust_stream_eligible: bool, + has_agentic_hook: bool, model: str, api_key: str | None, api_base: str | None, @@ -2279,7 +2277,7 @@ class BaseLLMHTTPHandler: return None if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled(): return None - if stream and not rust_stream_eligible: + if has_agentic_hook: return None from litellm.rust_bridge import messages as rust_messages_bridge diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b745ca8eadf..fbd7e36e298 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -219,8 +219,7 @@ def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), - "stream": False, - "rust_stream_eligible": False, + "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", "api_base": "https://resource.services.ai.azure.com/anthropic", @@ -345,11 +344,11 @@ async def test_gate_skips_rust_for_unsupported_provider(): @pytest.mark.asyncio -async def test_gate_skips_rust_when_streaming_but_not_eligible(): +async def test_gate_skips_rust_for_agentic_hook(): bridge = ExplodingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) - response = await _gate(stream=True, rust_stream_eligible=False) + response = await _gate(has_agentic_hook=True) assert response is None assert bridge.calls == 0 @@ -362,8 +361,7 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): streaming_body = {**REQUEST_BODY, "stream": True} response = await _gate( - stream=True, - rust_stream_eligible=True, + has_agentic_hook=False, request_body=streaming_body, ) From 06169e8c313bcb7e35591bdc96a1521338cee1d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:37:51 -0700 Subject: [PATCH 016/130] docs: add TLDR section to PR template --- .github/pull_request_template.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index d7e80b32749..5207e03fdc2 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,3 +1,8 @@ +## TLDR + + + ## Relevant issues From e7b9357bc2f913634ff073a81b9c26efec658313 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 21 Jul 2026 17:42:20 -0700 Subject: [PATCH 017/130] fix(e2e): drop httpbin.org from passthrough headers test, use real Anthropic (#34159) httpbin.org is an external dependency prone to transient 503s (caused the stage failure); its echo-body assertion also doesn't exercise a real LLM provider. Point the custom pass-through endpoint at the real Anthropic Messages API instead. Anthropic doesn't echo headers back, but it gates real behavior on two of them, which is enough to prove forwarding: a static x-api-key configured on the endpoint (never supplied by the caller) must reach upstream or the call 401s, and an invalid x-pass-anthropic-version sent by the caller must reach upstream with the prefix stripped, which Anthropic echoes verbatim in its 400 body. Verified live against a local proxy and the real Anthropic API: valid version returns a real completion, invalid version returns the exact marker in the 400 body. --- .../test_passthrough_headers_e2e.py | 92 ++++++++++--------- 1 file changed, 47 insertions(+), 45 deletions(-) diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index d9fe37c79ed..045988334d5 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -1,29 +1,32 @@ """Live e2e: custom pass-through endpoints inject configured headers and honor x-pass-* client headers (prefix stripped) on the way to the upstream. -The upstream is a real public echo service (httpbin.org/anything). Creating the -route via POST /config/pass_through_endpoint, calling it with a virtual key, and -asserting the echo body is the product path operators use; a mock would not -prove the proxy actually rewrote the outbound request. +The upstream is the real Anthropic Messages API rather than an echo service: +Anthropic doesn't echo request headers back, but it does gate real behavior on +two of them, which is enough to prove forwarding without a mock. A static +x-api-key configured on the pass-through endpoint (the caller never supplies +one) must reach upstream, or every call 401s; an invalid x-pass-anthropic-version +sent by the caller must reach upstream with the prefix stripped, and Anthropic +echoes the exact value back in its 400 body, so a unique-per-run marker proves +this specific request's header - not a stale or cached one - got there. """ from __future__ import annotations import pytest -from pydantic import BaseModel, Field, ValidationError +from pydantic import BaseModel, Field from e2e_config import unique_marker -from e2e_http import AuthHeaders, NoBody, StreamingResponse, require_successful_call, unwrap +from e2e_http import AuthHeaders, NoBody, require_successful_call, unwrap +from endpoints_client import MessagesResult from lifecycle import ResourceManager -from models import KeyGenerateBody +from models import ChatMessage, KeyGenerateBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e -ECHO_TARGET = "https://httpbin.org/anything" -STATIC_HEADER_NAME = "x-e2e-static-header" -PASS_HEADER_STEM = "e2e-client-marker" -PASS_HEADER_NAME = f"x-pass-{PASS_HEADER_STEM}" +ANTHROPIC_MESSAGES_TARGET = "https://api.anthropic.com/v1/messages" +MODEL = "claude-haiku-4-5-20251001" class PassThroughCreateBody(BaseModel): @@ -48,30 +51,26 @@ class PassThroughDeleteParams(BaseModel): endpoint_id: str -class EchoCallHeaders(AuthHeaders): +class AnthropicPassThroughHeaders(AuthHeaders): content_type: str = Field(default="application/json", serialization_alias="Content-Type") - x_pass_e2e_client_marker: str = Field(serialization_alias="x-pass-e2e-client-marker") + x_pass_anthropic_version: str = Field(serialization_alias="x-pass-anthropic-version") -class EchoBody(BaseModel): - ping: str +class AnthropicMessagesBody(BaseModel): + model: str + max_tokens: int = 8 + messages: list[ChatMessage] -class EchoResponse(BaseModel): - headers: dict[str, str] - - -def _create_passthrough( - client: PassthroughClient, *, path: str, static_value: str -) -> PassThroughEndpoint: +def _create_passthrough(client: PassthroughClient, *, path: str) -> PassThroughEndpoint: created = unwrap( client.proxy.transport.post( "/config/pass_through_endpoint", headers=client.proxy.transport.master, json=PassThroughCreateBody( path=path, - target=ECHO_TARGET, - headers={STATIC_HEADER_NAME: static_value}, + target=ANTHROPIC_MESSAGES_TARGET, + headers={"x-api-key": "os.environ/ANTHROPIC_API_KEY"}, ), response_type=PassThroughCreateResponse, ) @@ -92,12 +91,8 @@ def _delete_passthrough(client: PassthroughClient, endpoint_id: str) -> None: ) -def _echo_headers(resp: StreamingResponse) -> dict[str, str]: - try: - echo = EchoResponse.model_validate_json(resp.body) - except ValidationError as exc: - pytest.fail(f"echo upstream did not return a headers map: {exc}; body={resp.body[:300]}") - return {k.lower(): v for k, v in echo.headers.items()} +def _messages_body() -> AnthropicMessagesBody: + return AnthropicMessagesBody(model=MODEL, messages=[ChatMessage(role="user", content="Say hi.")]) class TestPassthroughHeaders: @@ -110,10 +105,8 @@ class TestPassthroughHeaders: ) -> None: marker = unique_marker() path = f"/e2e-passthrough-headers-{marker}" - static_value = f"static-{marker}" - client_value = f"client-{marker}" - endpoint = _create_passthrough(client, path=path, static_value=static_value) + endpoint = _create_passthrough(client, path=path) assert endpoint.id is not None resources.defer(lambda: _delete_passthrough(client, endpoint.id or "")) @@ -128,23 +121,32 @@ class TestPassthroughHeaders: result = client.proxy.transport.send( path, - headers=EchoCallHeaders( + headers=AnthropicPassThroughHeaders( authorization=f"Bearer {key}", - x_pass_e2e_client_marker=client_value, + x_pass_anthropic_version="2023-06-01", ), - json=EchoBody(ping=marker), + json=_messages_body(), ) require_successful_call(result) + completion = MessagesResult.model_validate_json(result.body) + assert completion.text.strip(), ( + f"static x-api-key must reach Anthropic for the call to succeed at all; got {result.body[:300]}" + ) - upstream = _echo_headers(result) - assert upstream.get(STATIC_HEADER_NAME) == static_value, ( - f"configured pass-through header {STATIC_HEADER_NAME!r} not on upstream " - f"request; got {upstream}" + invalid_version = f"e2e-passhdr-{unique_marker()}" + blocked = client.proxy.transport.send( + path, + headers=AnthropicPassThroughHeaders( + authorization=f"Bearer {key}", + x_pass_anthropic_version=invalid_version, + ), + json=_messages_body(), ) - assert upstream.get(PASS_HEADER_STEM) == client_value, ( - f"x-pass-* header should strip the prefix and forward as {PASS_HEADER_STEM!r}; " - f"got {upstream}" + assert blocked.status_code == 400, ( + f"expected Anthropic to reject the invalid anthropic-version, got " + f"{blocked.status_code}: {blocked.body[:300]}" ) - assert PASS_HEADER_NAME not in upstream, ( - "upstream must not see the x-pass- prefix; proxy should strip it" + assert invalid_version in blocked.body, ( + f"x-pass-anthropic-version must reach upstream with the prefix stripped; " + f"marker missing from Anthropic's error body: {blocked.body[:300]}" ) From 06d2efdd56061bca67ded17d195ca8d39132acff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:42:25 -0700 Subject: [PATCH 018/130] docs: structure PR template TLDR into problem/solution bullets --- .github/pull_request_template.md | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 5207e03fdc2..bbe5845f5a9 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,7 +1,17 @@ ## TLDR - + + +Problem this solves: + +- +- ... + +How it solves it: + +- +- ... ## Relevant issues From e967bc8c4f748e3302983ccdccce43bf7a33c5dc Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 17:43:41 -0700 Subject: [PATCH 019/130] test(e2e): cover 12 non-core LLM coverage registry cells (#34123) * fix(e2e): reference client.proxy in mid-conversation native providers test EndpointsClient exposes the shared ProxyClient as .proxy and has never had a .gateway attribute, so these two calls raised AttributeError at runtime and failed the tests/e2e basedpyright zero-error gate for any PR touching e2e files. Introduced in 23b5b7d199. * test(e2e): cover 12 non-core LLM coverage registry cells Raises Non-Core LLMs registry coverage from 24/50 to 36/50 (overall 51.9% to 54.8%). Four cells were already asserted by existing tests and only gain their covers marker (openai embeddings, openai image generation, openai TTS, cohere rerank); one is dual-marked onto the existing spend-tracking embeddings test rather than duplicated. New tests: bedrock and vertex embeddings, streaming TTS (asserts chunked transfer encoding so a buffered body cannot pass), audio transcriptions via the realtime suite's wav fixture, moderations flag/pass pair, and files list/retrieve in the batches suite. Harness: e2e_http.upload generalized to any form model with a file_content_type override (batches path unchanged), new stream_binary primitive + BinaryStream for binary chunked responses, transcribe and moderations client methods, file retrieve/list client methods. * fix(e2e): close streamed TTS response on error paths and surface the error body With stream=True a non-2xx response returned with the body unread, keeping the socket checked out until garbage collection; the sibling _streaming_outcome already consumes resp.text on error. The response now closes on every path and BinaryStream carries a bounded error_body so a failed stream call is triageable. * test(e2e): assert streamed TTS response carries no content-length --- tests/e2e/batches/batch_client.py | 28 +++++- tests/e2e/batches/test_batches_e2e.py | 67 +++++++++++++++ tests/e2e/e2e_http.py | 86 +++++++++++++++++-- tests/e2e/llm_translation/endpoints_client.py | 63 +++++++++++++- .../llm_translation/test_audio_speech_e2e.py | 46 +++++++++- .../test_audio_transcriptions_e2e.py | 51 +++++++++++ .../test_embeddings_endpoint_e2e.py | 57 ++++++++++-- .../test_image_generation_e2e.py | 1 + .../llm_translation/test_moderations_e2e.py | 65 ++++++++++++++ tests/e2e/llm_translation/test_rerank_e2e.py | 1 + .../spend_tracking/test_spend_tracking_e2e.py | 1 + tests/e2e/transport.py | 50 ++++++++++- 12 files changed, 496 insertions(+), 20 deletions(-) create mode 100644 tests/e2e/llm_translation/test_audio_transcriptions_e2e.py create mode 100644 tests/e2e/llm_translation/test_moderations_e2e.py diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 7db5d0b6beb..5cc5d1dae3b 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -26,16 +26,24 @@ from e2e_http import ( ) from models import LiteLLMParamsBody +UPLOAD_FILENAME = "batch_input.jsonl" + class FileObject(BaseModel): id: str object: str | None = None purpose: str | None = None + filename: str | None = None bytes: int | None = None status: str | None = None created_at: int | None = None +class FileList(BaseModel): + object: str | None = None + data: list[FileObject] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -106,12 +114,30 @@ class BatchClient: _files_path(provider), headers=self.proxy.transport.bearer(key), form=form, - filename="batch_input.jsonl", + filename=UPLOAD_FILENAME, content=content, params=ModelQuery(model=model), response_type=FileObject, ) + def retrieve_file( + self, file_id: str, *, key: str, provider: str | None = None + ) -> Result[FileObject]: + return self.proxy.transport.get( + f"{_files_path(provider)}/{file_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=FileObject, + ) + + def list_files(self, *, key: str, provider: str | None = None) -> Result[FileList]: + return self.proxy.transport.get( + _files_path(provider), + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=FileList, + ) + def create_batch( self, *, body: BatchCreateBody, key: str, provider: str | None = None ) -> StreamingResponse: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index f886c09b705..b0c53becb6b 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -26,6 +26,7 @@ import pytest from e2e_config import require_env, unique_marker from batch_client import ( + UPLOAD_FILENAME, BatchClient, BatchCreateBody, BatchObject, @@ -511,6 +512,72 @@ class TestBatchFileContent: ) +class TestOpenAIFiles: + """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. + + The proxy lists the OpenAI org's raw file ids, so the list case uploads a raw + (provider-routed) file whose id matches what list returns; retrieve re-encodes + the id it was called with, so the model-encoded upload round-trips unchanged. + """ + + @pytest.mark.covers( + "llm.files.openai.list.nonstream.works", + exercised_on=["files"], + ) + def test_uploaded_file_appears_in_list( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(OPENAI_BATCH_MODEL), + form=FileUploadForm(purpose="batch"), + key=key, + provider="openai", + ) + ) + resources.defer( + quietly(lambda: client.delete_file(file.id, key=key, provider="openai")) + ) + + listed = unwrap(client.list_files(key=key)) + assert listed.object is None or listed.object == "list", ( + f"list envelope object={listed.object!r}" + ) + match = next((entry for entry in listed.data if entry.id == file.id), None) + assert match is not None, f"uploaded file {file.id!r} absent from GET /v1/files" + assert match.purpose == "batch", ( + f"listed file must round-trip the upload purpose, got {match.purpose!r}" + ) + + @pytest.mark.covers( + "llm.files.openai.retrieve.nonstream.works", + exercised_on=["files"], + ) + def test_retrieve_round_trips_metadata( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(OPENAI_BATCH_MODEL), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + fetched = unwrap(client.retrieve_file(file.id, key=key)) + assert fetched.id == file.id, "retrieve must echo the uploaded file id" + assert fetched.purpose == "batch", ( + f"retrieve must round-trip purpose, got {fetched.purpose!r}" + ) + assert fetched.filename == UPLOAD_FILENAME, ( + f"retrieve must round-trip filename, got {fetched.filename!r}" + ) + + BATCH_RL_REQUEST_LINES = 3 BATCH_RL_RPM_LIMIT = 2 diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 7ec4a439e10..2f54c789c6b 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -147,6 +147,32 @@ class StreamingResponse(BaseModel): return "text/event-stream" in (self.content_type or "") +class BinaryStream(BaseModel): + """Outcome of consuming a binary chunked response (e.g. TTS audio) as a stream. + + Unlike StreamingResponse, which line-splits an SSE text body, this iterates the + raw bytes with iter_content and reports how many non-empty chunks arrived and + the total byte count, so a caller can assert customer-observable streaming + (multiple chunks, real bytes) without decoding the payload.""" + + status_code: int + content_type: str | None = None + call_id: str | None = None + transfer_encoding: str | None = None + content_length: str | None = None + error_body: str | None = None + chunk_count: int = 0 + total_bytes: int = 0 + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + @property + def chunked(self) -> bool: + return "chunked" in (self.transfer_encoding or "") + + def _hdr(resp: requests.Response, name: str) -> str | None: value = resp.headers.get(name) return value if isinstance(value, str) else None @@ -430,16 +456,18 @@ def upload[R: BaseModel]( url: URL, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], timeout: float = 60.0, ) -> Result[R]: - """Multipart POST for file uploads (/v1/files). Form fields come from `form`, - the file bytes are sent as the `file` part, and `params` carries any query - routing (e.g. ?model=). requests sets the multipart Content-Type itself.""" + """Multipart POST for file-bearing routes (/v1/files, /v1/audio/transcriptions). + Form fields come from `form`, the file bytes are sent as the `file` part with + `file_content_type`, and `params` carries any query routing (e.g. ?model=). + requests sets the multipart Content-Type itself.""" dumped: dict[str, object] = form.model_dump(by_alias=True, exclude_none=True) data = {key: str(value) for key, value in dumped.items()} try: @@ -448,7 +476,7 @@ def upload[R: BaseModel]( headers=_headers(headers), params=_params(params), data=data, - files={"file": (filename, content, "application/jsonl")}, + files={"file": (filename, content, file_content_type)}, timeout=timeout, ) except requests.RequestException as exc: @@ -456,6 +484,54 @@ def upload[R: BaseModel]( return _classify(resp, response_type) +def stream_binary( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + timeout: float = 60.0, +) -> BinaryStream: + """POST that consumes a binary chunked response (e.g. TTS audio) as a stream, + counting non-empty chunks and total bytes with iter_content. A non-2xx status + short-circuits with the counts left at zero so the caller can fail loudly.""" + try: + resp = requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + stream=True, + timeout=timeout, + ) + except requests.RequestException as exc: + return BinaryStream(status_code=-1, error_body=str(exc)[:300]) + with resp: + content_type = _hdr(resp, "content-type") + call_id = _hdr(resp, "x-litellm-call-id") + transfer_encoding = _hdr(resp, "transfer-encoding") + content_length = _hdr(resp, "content-length") + if not (200 <= resp.status_code < 300): + return BinaryStream( + status_code=resp.status_code, + content_type=content_type, + call_id=call_id, + transfer_encoding=transfer_encoding, + content_length=content_length, + error_body=resp.text[:300], + ) + raw_chunks = cast("Iterator[bytes]", resp.iter_content(chunk_size=chunk_size)) + chunks = tuple(chunk for chunk in raw_chunks if chunk) + return BinaryStream( + status_code=resp.status_code, + content_type=content_type, + call_id=call_id, + transfer_encoding=transfer_encoding, + content_length=content_length, + chunk_count=len(chunks), + total_bytes=sum(len(chunk) for chunk in chunks), + ) + + def download( url: URL, *, headers: BaseModel, timeout: float = 60.0 ) -> StreamingResponse: diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 32d9922c775..ace621d03b3 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -15,7 +15,7 @@ from typing import Literal from pydantic import BaseModel from proxy_client import ProxyClient -from e2e_http import StreamingResponse +from e2e_http import BinaryStream, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock __all__ = [ @@ -110,6 +110,16 @@ class ImageRequest(BaseModel): size: str = "1024x1024" +class TranscriptionForm(BaseModel): + model: str + response_format: str = "json" + + +class ModerationRequest(BaseModel): + model: str + input: str + + class ResponsesOutputContent(BaseModel): type: str | None = None text: str | None = None @@ -213,6 +223,27 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] +class TranscriptionResult(BaseModel): + text: str = "" + + +class ModerationResultItem(BaseModel): + flagged: bool + categories: dict[str, bool] = {} + + @property + def flagged_categories(self) -> tuple[str, ...]: + return tuple(name for name, hit in self.categories.items() if hit) + + +class ModerationResult(BaseModel): + results: list[ModerationResultItem] = [] + + @property + def first(self) -> ModerationResultItem | None: + return self.results[0] if self.results else None + + @dataclass(frozen=True, slots=True) class EndpointsClient: proxy: ProxyClient @@ -314,6 +345,36 @@ class EndpointsClient: "/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice) ) + def audio_speech_stream( + self, key: str, model: str, text: str, *, voice: str = "alloy" + ) -> BinaryStream: + return self.proxy.transport.stream_binary( + "/v1/audio/speech", + headers=self.proxy.transport.bearer(key), + json=SpeechRequest(model=model, input=text, voice=voice), + ) + + def transcribe( + self, key: str, model: str, *, filename: str, content: bytes + ) -> Result[TranscriptionResult]: + return self.proxy.transport.upload( + "/v1/audio/transcriptions", + headers=self.proxy.transport.bearer(key), + form=TranscriptionForm(model=model), + filename=filename, + content=content, + file_content_type="audio/wav", + response_type=TranscriptionResult, + ) + + def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]: + return self.proxy.transport.post( + "/v1/moderations", + headers=self.proxy.transport.bearer(key), + json=ModerationRequest(model=model, input=text), + response_type=ModerationResult, + ) + def images(self, key: str, model: str, prompt: str) -> StreamingResponse: return self._send( "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) diff --git a/tests/e2e/llm_translation/test_audio_speech_e2e.py b/tests/e2e/llm_translation/test_audio_speech_e2e.py index f7a04d94cb3..b95cef8db4d 100644 --- a/tests/e2e/llm_translation/test_audio_speech_e2e.py +++ b/tests/e2e/llm_translation/test_audio_speech_e2e.py @@ -1,8 +1,9 @@ -"""Live e2e: POST /v1/audio/speech returns audio. +"""Live e2e: POST /v1/audio/speech returns audio, non-streamed and streamed. -Registers an OpenAI text-to-speech deployment at runtime and asserts the response -is an audio body (binary, not JSON). Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +The non-streamed call asserts an audio (not JSON) body. The streamed call consumes +the response the way a player would and asserts customer-observable streaming: +chunked transfer encoding (a buffered body would carry a content-length) with +non-zero audio bytes. """ from __future__ import annotations @@ -19,6 +20,7 @@ pytestmark = pytest.mark.e2e class TestAudioSpeech: + @pytest.mark.covers("llm.audio_speech.openai.basic.nonstream.works") def test_audio_speech_returns_audio( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -38,3 +40,39 @@ class TestAudioSpeech: f"/audio/speech content-type is not audio: {result.content_type!r}" ) assert result.body, "/audio/speech returned an empty body" + + @pytest.mark.covers("llm.audio_speech.openai.basic.stream.works") + def test_audio_speech_streams_audio_chunks( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-speech-stream-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.audio_speech_stream( + key, + model, + "Streaming speech should arrive in several audio chunks so a client can " + "begin playback well before the whole clip has finished generating.", + ) + assert result.ok, ( + f"/audio/speech stream failed (status {result.status_code}); body={result.error_body}" + ) + assert "audio" in (result.content_type or ""), ( + f"/audio/speech content-type is not audio: {result.content_type!r}" + ) + assert result.chunked, ( + f"/audio/speech did not stream: transfer-encoding={result.transfer_encoding!r}, " + f"content-length={result.content_length!r} (a buffered body is not a stream)" + ) + assert result.content_length is None, ( + f"/audio/speech advertised content-length={result.content_length!r} on a " + f"streamed response (a buffered body is not a stream)" + ) + assert result.total_bytes > 0, "/audio/speech stream returned no audio bytes" diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py new file mode 100644 index 00000000000..af6123dc46a --- /dev/null +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -0,0 +1,51 @@ +"""Live e2e: POST /v1/audio/transcriptions turns speech into text. + +Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken +weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting +the returned transcript is non-empty and mentions the word it was asked about. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +WEATHER_WAV = ( + Path(__file__).resolve().parent / "realtime" / "fixtures" / "weather_question_24k.wav" +) + + +class TestAudioTranscriptions: + @pytest.mark.covers("llm.audio_transcriptions.openai.basic.nonstream.works") + def test_audio_transcriptions_returns_text( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-transcribe-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-4o-mini-transcribe", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = unwrap( + endpoints_client.transcribe( + key, model, filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes() + ) + ) + text = result.text.strip() + assert text, "/audio/transcriptions returned an empty transcript" + assert "weather" in text.lower(), ( + f"transcript of a spoken weather question does not mention weather: {text!r}" + ) diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 56f2de8bd4f..157caedd561 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,9 +1,9 @@ -"""Live e2e: POST /embeddings returns a real vector. +"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. -Registers an OpenAI embedding deployment at runtime and asserts a non-empty, -non-zero vector came back. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py; the LIT-3167 guard in -tests/e2e/embeddings/ covers the Gemini embedding path. +Each test registers the deployment it needs at runtime (deleted on teardown) and +asserts a non-empty, non-zero vector came back. The LIT-3167 guard in +tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking is +covered by tests/e2e/quota_management/spend_tracking/. """ from __future__ import annotations @@ -20,6 +20,7 @@ pytestmark = pytest.mark.e2e class TestEmbeddingsEndpoint: + @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: @@ -40,3 +41,49 @@ class TestEmbeddingsEndpoint: assert any(component != 0.0 for component in parsed.first_vector), ( f"embedding vector is all zeros: {result.body[:300]}" ) + + @pytest.mark.covers("llm.embeddings.bedrock.basic.nonstream.works") + def test_bedrock_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-bedrock-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.titan-embed-text-v2:0", aws_region_name="us-west-2" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) + + @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") + def test_vertex_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-vertex-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="vertex_ai/gemini-embedding-2", + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="us-central1", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 4d2211f3be4..d4080afb7dc 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -19,6 +19,7 @@ pytestmark = pytest.mark.e2e class TestImageGeneration: + @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/llm_translation/test_moderations_e2e.py b/tests/e2e/llm_translation/test_moderations_e2e.py new file mode 100644 index 00000000000..69cf4414a48 --- /dev/null +++ b/tests/e2e/llm_translation/test_moderations_e2e.py @@ -0,0 +1,65 @@ +"""Live e2e: POST /v1/moderations classifies content against the provider policy. + +Registers OpenAI's omni moderation model at runtime and asserts the product +promise on both sides of the decision: clearly violent text comes back flagged +with at least one policy category tripped, and benign text comes back not flagged. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +VIOLENT_TEXT = "I am going to find you and kill you, and I will hurt everyone you love." +BENIGN_TEXT = "I enjoyed the sunny afternoon and a relaxing walk in the park today." + + +def _register_moderation_model( + endpoints_client: EndpointsClient, resources: ResourceManager +) -> str: + model = f"e2e-moderation-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/omni-moderation-latest", api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + return model + + +class TestModerations: + @pytest.mark.covers("llm.moderations.openai.basic.nonstream.works") + def test_moderations_flags_violent_content( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_moderation_model(endpoints_client, resources) + key = resources.key() + + result = unwrap(endpoints_client.moderations(key, model, VIOLENT_TEXT)) + item = result.first + assert item is not None, f"/moderations returned no results: {result}" + assert item.flagged, f"violent text was not flagged: {item}" + assert item.flagged_categories, ( + f"flagged result reported no true category: {item}" + ) + + def test_moderations_passes_benign_content( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_moderation_model(endpoints_client, resources) + key = resources.key() + + result = unwrap(endpoints_client.moderations(key, model, BENIGN_TEXT)) + item = result.first + assert item is not None, f"/moderations returned no results: {result}" + assert not item.flagged, ( + f"benign text was flagged as {item.flagged_categories}: {item}" + ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 4b30ac1ea5c..31801b306a8 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -26,6 +26,7 @@ DOCUMENTS = [ class TestRerank: + @pytest.mark.covers("llm.rerank.cohere.basic.nonstream.works") def test_rerank_scores_top_n( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index d43d8e94898..3bd1b992745 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -194,6 +194,7 @@ def test_streaming_messages_via_responses_bridge_tracks_spend( @pytest.mark.covers("quota_management.spend_tracking.embeddings.logs_cost") +@pytest.mark.covers("llm.embeddings.openai.basic.nonstream.cost_logged") def test_embedding_writes_nonzero_spend_row( client: SpendClient, scoped_key: str ) -> None: diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index f7061b03a46..da4252e550e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -16,7 +16,7 @@ import e2e_http from e2e_http import ( URL, AuthHeaders, - FileUploadForm, + BinaryStream, ProbeResult, Result, StreamingResponse, @@ -32,6 +32,15 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel ) -> StreamingResponse: ... + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: ... + def send( self, path: str, @@ -76,9 +85,10 @@ class Transport(Protocol): path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: ... @@ -181,6 +191,22 @@ class HttpTransport: self._url(path), headers=headers, json=json, timeout=self.request_timeout ) + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: + return e2e_http.stream_binary( + self._url(path), + headers=headers, + json=json, + chunk_size=chunk_size, + timeout=self.request_timeout, + ) + def send( self, path: str, @@ -212,9 +238,10 @@ class HttpTransport: path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: @@ -224,6 +251,7 @@ class HttpTransport: form=form, filename=filename, content=content, + file_content_type=file_content_type, params=params, response_type=response_type, timeout=self.request_timeout, @@ -346,6 +374,18 @@ class SplitTransport: ) -> StreamingResponse: return self._route(path).stream(path, headers=headers, json=json) + def stream_binary( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + chunk_size: int = 8192, + ) -> BinaryStream: + return self._route(path).stream_binary( + path, headers=headers, json=json, chunk_size=chunk_size + ) + def send( self, path: str, @@ -367,9 +407,10 @@ class SplitTransport: path: str, *, headers: BaseModel, - form: FileUploadForm, + form: BaseModel, filename: str, content: bytes, + file_content_type: str = "application/jsonl", params: BaseModel | None = None, response_type: type[R], ) -> Result[R]: @@ -379,6 +420,7 @@ class SplitTransport: form=form, filename=filename, content=content, + file_content_type=file_content_type, params=params, response_type=response_type, ) From 4c1f071adde23df1bb5379f674fb490e8977226a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:44:02 -0700 Subject: [PATCH 020/130] docs: cap PR template TLDR bullets at one short line --- .github/pull_request_template.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index bbe5845f5a9..1301bfb0e60 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,6 +1,6 @@ ## TLDR - Problem this solves: From 5081e0cf797a5999807079022dc7a48ad21e53c6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 17:56:20 -0700 Subject: [PATCH 021/130] test(logging): pin compression_savings in the gcs pubsub spend log fixture (#34204) The spend-log metadata schema gained a compression_savings key, so the gcs pubsub v1 payload now carries it. The golden fixture was never updated, and the comparator flags any key present in the payload but absent from the fixture, so test_async_gcs_pub_sub_v1 failed on every run. Pin the key as null rather than adding it to ignored_keys; the value is deterministic on this path, so ignoring it would leave the assertion blind to the field entirely. --- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index a4c50d3c575..41b4c3efb63 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From 065faf6e695be64a1a367601271839914e4a6621 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 21 Jul 2026 17:57:58 -0700 Subject: [PATCH 022/130] chore(proxy): clean up request parameter validation and provider destination handling (#34189) --- litellm/litellm_core_utils/url_utils.py | 9 + litellm/llms/huggingface/embedding/handler.py | 2 +- .../huggingface/embedding/transformation.py | 19 - litellm/llms/oobabooga/chat/oobabooga.py | 4 +- litellm/proxy/auth/auth_utils.py | 72 +++- litellm/proxy/auth/user_api_key_auth.py | 59 +-- litellm/proxy/litellm_pre_call_utils.py | 40 +- .../code_coverage_tests/recursive_detector.py | 1 + .../test_huggingface_embedding_handler.py | 14 + .../llms/oobabooga/chat/test_oobabooga.py | 55 +++ .../proxy/auth/test_auth_utils.py | 352 +++++++++++++++++- .../test_router_override_fallback_auth.py | 141 ++++++- .../test_provider_url_destination_guard.py | 40 ++ 13 files changed, 702 insertions(+), 106 deletions(-) create mode 100644 tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1cbb1ce973f..a83cb3bc69e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -148,6 +148,15 @@ def _parse_url_destination_allowlist_entry( return _normalize_host(parsed.hostname), scheme, port +def provider_url_destination_candidates(value: str) -> Tuple[str, ...]: + return tuple( + candidate + for part in value.split(",") + for candidate in (part.strip(), part.strip().split("/", 1)[1] if "/" in part.strip() else "") + if candidate + ) + + def is_url_destination_allowed_by_host(url: str, allowed_hosts: List[str]) -> bool: """Return True when a credential-bearing provider URL is admin-allowlisted. diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 39eb430db74..f72a79e084d 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -322,7 +322,7 @@ class HuggingFaceEmbedding(BaseLLM): task = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) # print_verbose(f"{model}, {task}") embed_url = "" - if "https" in model: + if model.startswith(("http://", "https://")): embed_url = model elif api_base: embed_url = api_base diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 13e38ab5560..6f27e3115eb 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -316,25 +316,6 @@ class HuggingFaceEmbeddingConfig(BaseConfig): return data - def get_api_base(self, api_base: Optional[str], model: str) -> str: - """ - Get the API base for the Huggingface API. - - Do not add the chat/embedding/rerank extension here. Let the handler do this. - """ - if "https" in model: - completion_url = model - elif api_base is not None: - completion_url = api_base - elif "HF_API_BASE" in os.environ: - completion_url = os.getenv("HF_API_BASE", "") - elif "HUGGINGFACE_API_BASE" in os.environ: - completion_url = os.getenv("HUGGINGFACE_API_BASE", "") - else: - completion_url = f"https://api-inference.huggingface.co/models/{model}" - - return completion_url - def validate_environment( self, headers: Dict, diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index fe2bb9dc6d1..40d88e8e125 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -34,7 +34,7 @@ def completion( optional_params=optional_params, litellm_params=litellm_params, ) - if "https" in model: + if model.startswith(("http://", "https://")): completion_url = model elif api_base: completion_url = api_base @@ -96,7 +96,7 @@ def embedding( encoding=None, ): # Create completion URL - if "https" in model: + if model.startswith(("http://", "https://")): embeddings_url = model elif api_base: embeddings_url = f"{api_base}/v1/embeddings" diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 293bb74e211..ecb37e67c14 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -3,7 +3,7 @@ import re import sys from functools import lru_cache from logging import Logger -from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Tuple, Union +from typing import Any, Dict, FrozenSet, Iterator, List, Mapping, Optional, Tuple, Union from fastapi import HTTPException, Request, status @@ -12,7 +12,12 @@ from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.litellm_core_utils.url_utils import ( + SSRFError, + is_url_destination_allowed_by_host, + provider_url_destination_candidates, + validate_url, +) from litellm.proxy._types import * from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, @@ -290,6 +295,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( "use_ssl", # SDK-only field; also rejected outright in is_request_body_safe. "model_list", + "vertex_ai_credentials", # Observability credentials, hosts, and project identifiers: derived # from the canonical ``_supported_callback_params`` allowlist so new # integrations are covered automatically. Sorted for stable iteration @@ -342,6 +348,60 @@ def _check_banned_params( ) +_FALLBACK_FIELDS: tuple[str, ...] = ( + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", +) + + +def _iter_fallback_field_values(request_body: Mapping[str, object]) -> Iterator[object]: + override = request_body.get("router_settings_override") + for source in (request_body, override): + if isinstance(source, Mapping): + for field in _FALLBACK_FIELDS: + yield source.get(field) + + +def _iter_fallback_targets(value: object, depth: int) -> Iterator[str | Mapping[str, object]]: + if depth > 2 * litellm.ROUTER_MAX_FALLBACKS: + raise ValueError("Rejected Request: fallback nesting exceeds the allowed validation depth.") + if not isinstance(value, list): + return + for item in value: + if isinstance(item, str): + yield item + elif isinstance(item, Mapping): + values = tuple(item.values()) + if not (values and all(isinstance(v, list) for v in values)): + yield item + if isinstance(item.get("model"), str): + for field in _FALLBACK_FIELDS: + yield from _iter_fallback_targets(item.get(field), depth + 1) + else: + for target_list in values: + yield from _iter_fallback_targets(target_list, depth + 1) + + +def iter_request_fallback_targets(request_body: Mapping[str, object]) -> Iterator[str | Mapping[str, object]]: + for value in _iter_fallback_field_values(request_body): + yield from _iter_fallback_targets(value, 0) + + +def _reject_url_valued_fallback_target(value: str) -> None: + allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] + for candidate in provider_url_destination_candidates(value): + if not candidate.lower().startswith(("http://", "https://")): + continue + if is_url_destination_allowed_by_host(candidate, allowed_hosts): + continue + raise ValueError( + f"Rejected Request: URL-valued fallback destination '{value}' is not allowed. " + "Configure custom endpoints with api_base instead, or add the destination host to " + "`provider_url_destination_allowed_hosts` in litellm_settings." + ) + + def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str) -> bool: """ Check if the request body is safe. @@ -379,6 +439,14 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) if metadata is not None: _check_banned_params(metadata, general_settings, llm_router, model) + for target in iter_request_fallback_targets(request_body): + if isinstance(target, dict): + _check_banned_params(target, general_settings, llm_router, model) + target_model = target.get("model") + if isinstance(target_model, str): + _reject_url_valued_fallback_target(target_model) + elif isinstance(target, str): + _reject_url_valued_fallback_target(target) litellm_params = _coerce_metadata_to_dict(request_body.get("litellm_params")) if litellm_params is not None: litellm_params_metadata = _coerce_metadata_to_dict(litellm_params.get("metadata")) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c185f950395..83a8a69511b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -14,7 +14,7 @@ import secrets import orjson from datetime import datetime, timezone -from typing import Any, Dict, Iterator, NamedTuple, List, Optional, Protocol, Tuple, Union, cast +from typing import Any, Dict, NamedTuple, List, Optional, Protocol, Tuple, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_utils import ( get_model_from_request, get_request_route, get_request_route_template, + iter_request_fallback_targets, normalize_request_route, pre_db_read_auth_checks, route_in_additonal_public_routes, @@ -2801,19 +2802,11 @@ async def _enforce_key_and_fallback_model_access( llm_router=llm_router, ) - # Validate every fallback model name reachable by this request. - # All three fields (``fallbacks``, ``context_window_fallbacks``, - # ``content_policy_fallbacks``) are forwarded to the router as - # per-request kwargs whether they appear at the top level of - # ``request_data`` or nested under ``router_settings_override``. - # Both surfaces must be validated against the API key's model - # allowlist or a caller can smuggle a restricted model. VERIA-44. - fallback_names: List[str] = [] - override_settings = request_data.get("router_settings_override") - for _fb_key in ROUTER_FALLBACK_FIELDS: - fallback_names.extend(iter_router_fallback_model_names(request_data.get(_fb_key))) - if isinstance(override_settings, dict): - fallback_names.extend(iter_router_fallback_model_names(override_settings.get(_fb_key))) + fallback_names = tuple( + name + for target in iter_request_fallback_targets(request_data) + if (name := _fallback_target_model_name(target)) is not None + ) for _name in dict.fromkeys(fallback_names): # dedupe, preserve order await can_key_call_model( @@ -2829,36 +2822,14 @@ async def _enforce_key_and_fallback_model_access( ) -ROUTER_FALLBACK_FIELDS: Tuple[str, ...] = ( - "fallbacks", - "context_window_fallbacks", - "content_policy_fallbacks", -) - - -def iter_router_fallback_model_names(fallbacks: Any) -> Iterator[str]: - """Yield leaf model names from any of the supported fallbacks shapes. - - Handles the simple top-level shape (``str`` or ``{"model": str}``) and - the nested router-config shape (``[{primary: [fallback_list]}]``). - """ - if not isinstance(fallbacks, list): - return - for entry in fallbacks: - if isinstance(entry, str): - yield entry - elif isinstance(entry, dict): - if isinstance(entry.get("model"), str): - yield entry["model"] - continue - for fallback_list in entry.values(): - if not isinstance(fallback_list, list): - continue - for m in fallback_list: - if isinstance(m, str): - yield m - elif isinstance(m, dict) and isinstance(m.get("model"), str): - yield m["model"] +def _fallback_target_model_name(target: object) -> str | None: + if isinstance(target, str): + return target + if isinstance(target, dict): + model = target.get("model") + if isinstance(model, str): + return model + return None async def _run_post_custom_auth_checks( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6514d4e1e8c..9d9ef28ec9b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -19,7 +19,10 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( iter_client_callback_metadata_dicts, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host +from litellm.litellm_core_utils.url_utils import ( + is_url_destination_allowed_by_host, + provider_url_destination_candidates, +) from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, @@ -227,23 +230,26 @@ def _reject_url_valued_destinations(data: Dict[str, Any]) -> None: allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] for field in _URL_DESTINATION_REQUEST_FIELDS: value = data.get(field) - if not isinstance(value, str) or not value.startswith(("http://", "https://")): + if not isinstance(value, str): continue - if is_url_destination_allowed_by_host(value, allowed_hosts): - continue - raise HTTPException( - status_code=400, - detail={ - "error": "invalid_request", - "param": field, - "message": ( - f"URL-valued '{field}' is not allowed. Configure custom " - "endpoints with api_base instead, or add the destination " - "host to `provider_url_destination_allowed_hosts` in " - "litellm_settings." - ), - }, - ) + for candidate in provider_url_destination_candidates(value): + if not candidate.lower().startswith(("http://", "https://")): + continue + if is_url_destination_allowed_by_host(candidate, allowed_hosts): + continue + raise HTTPException( + status_code=400, + detail={ + "error": "invalid_request", + "param": field, + "message": ( + f"URL-valued '{field}' is not allowed. Configure custom " + "endpoints with api_base instead, or add the destination " + "host to `provider_url_destination_allowed_hosts` in " + "litellm_settings." + ), + }, + ) def _strip_untrusted_request_header_controls( diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index fa81efde5db..0bc3cebdd5a 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -56,6 +56,7 @@ IGNORE_FUNCTIONS = [ "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_filter_mcp_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the MCP call at the cap. "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. + "_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap. ] diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 8a072fa5097..af8321f24a1 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -121,6 +121,20 @@ class TestHuggingFaceEmbedding: assert response.usage.prompt_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + def test_model_name_with_https_substring_uses_api_base(self): + api_base = "https://legit.example/embed" + + litellm.embedding( + model="huggingface/my-https-endpoint", + input=["hello world"], + input_type="embed", + api_base=api_base, + ) + + self.mock_http.assert_called_once() + called_url = self.mock_http.call_args[0][0] + assert called_url == api_base + def test_embedding_with_sentence_similarity_task(self): """Test embedding when task type is sentence-similarity (requires 2+ sentences)""" diff --git a/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py b/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py new file mode 100644 index 00000000000..91ebb2bd9d4 --- /dev/null +++ b/tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py @@ -0,0 +1,55 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm + +MOCK_COMPLETION_RESPONSE = { + "choices": [{"message": {"role": "assistant", "content": "hi there"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, +} + + +def _mock_post_response(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = "ok" + mock_response.json.return_value = MOCK_COMPLETION_RESPONSE + return mock_response + + +def test_model_name_with_https_substring_uses_api_base(): + api_base = "https://legit.example" + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" + ) as mock_post: + mock_post.return_value = _mock_post_response() + + litellm.completion( + model="oobabooga/my-https-model", + messages=[{"role": "user", "content": "hello"}], + api_base=api_base, + ) + + mock_post.assert_called_once() + called_url = mock_post.call_args[0][0] + assert called_url == f"{api_base}/v1/chat/completions" + + +def test_url_valued_model_still_targets_that_url(): + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" + ) as mock_post: + mock_post.return_value = _mock_post_response() + + litellm.completion( + model="oobabooga/https://sdk-user.example", + messages=[{"role": "user", "content": "hello"}], + ) + + mock_post.assert_called_once() + called_url = mock_post.call_args[0][0] + assert called_url == "https://sdk-user.example/v1/chat/completions" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 72bd215b9be..9f24c662581 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1587,7 +1587,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"base_url": "https://attacker.example"}, + request_kwargs={"base_url": "https://attacker.example", "api_key": "sk-caller"}, ) assert "aws_access_key_id" not in out assert "aws_secret_access_key" not in out @@ -1608,7 +1608,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"api_base": "self-hosted.example.com:50051"}, + request_kwargs={"api_base": "self-hosted.example.com:50051", "api_key": "sk-caller"}, ) assert out["api_base"] == "self-hosted.example.com:50051" assert "nvcf_function_id" not in out @@ -1626,7 +1626,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: } out = get_dynamic_litellm_params( litellm_params=dict(admin_params), - request_kwargs={"api_base": "self-hosted.example.com:50051"}, + request_kwargs={"api_base": "self-hosted.example.com:50051", "api_key": "sk-caller"}, ) assert out["api_base"] == "self-hosted.example.com:50051" assert "use_ssl" not in out @@ -1651,6 +1651,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: }, request_kwargs={ "api_base": "https://attacker.example", + "api_key": "sk-caller", "organization": "org-attacker", "extra_body": {"attacker": "value"}, }, @@ -1674,6 +1675,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: }, request_kwargs={ "api_base": "https://attacker.example", + "api_key": "sk-caller", "organization": "", "extra_body": "", }, @@ -1701,6 +1703,310 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert out["api_version"] == "2026-04-01" assert out["api_base"] == "https://admin.upstream/v1" + def test_client_api_key_used_when_supplied_with_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "model": "gpt-4", + "api_key": "sk-admin-secret", + "api_base": "https://admin.upstream/v1", + }, + request_kwargs={ + "api_base": "https://attacker.example", + "api_key": "sk-client-byok", + }, + ) + assert out["api_key"] == "sk-client-byok" + assert "sk-admin-secret" not in str(out) + + +_OPENAI_CHAT_RESPONSE = { + "id": "chatcmpl-x", + "object": "chat.completion", + "created": 1, + "model": "gpt-4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class TestClientsideBaseOverrideOutboundKey: + """Drive a completion through the router and assert on the outbound request + when the caller overrides ``api_base``.""" + + def _router(self): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-SERVER-CONFIG", + "api_base": "https://admin.upstream/v1", + }, + } + ] + ) + + @pytest.fixture(autouse=True) + def _ambient_server_key(self, monkeypatch): + import litellm + + monkeypatch.setenv("OPENAI_API_KEY", "sk-SERVER-ENV") + monkeypatch.setattr(litellm, "api_key", None, raising=False) + + def test_caller_key_override_sends_caller_key_never_server_key(self): + import httpx + import respx + + with respx.mock: + route = respx.post("https://caller.example/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_OPENAI_CHAT_RESPONSE) + ) + self._router().completion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://caller.example/v1", + api_key="sk-CALLER", + ) + authorization = route.calls.last.request.headers.get("authorization") + assert authorization == "Bearer sk-CALLER" + assert "SERVER" not in (authorization or "") + + +def _rounds_deep_api_base_payload(rounds, field): + """Build a fallbacks payload with ``api_base`` on a target nested ``rounds`` + fallback-rounds deep, each round wrapped in its own grouping dict.""" + node = {"model": "leaf", "api_base": "https://attacker.example"} + for i in range(rounds): + node = {"model": f"m{i}", field: [{"grp": [node]}]} + return {"model": "gpt-4", field: [{"grp": [node]}]} + + +class TestIsRequestBodySafeBlocksFallbackSmuggle: + """``is_request_body_safe`` runs the banned-param check on every dict target + inside the fallback lists.""" + + @pytest.fixture(autouse=True) + def _disable_url_validation(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + + @pytest.mark.parametrize( + "fallback_key", + ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"], + ) + def test_api_base_smuggled_via_nested_fallback_is_rejected(self, fallback_key): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_key: [ + { + "gpt-4": [ + {"model": "evil", "api_base": "https://attacker.example"}, + ] + } + ], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_string_only_fallbacks_are_accepted(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_benign_dict_fallback_entry_is_accepted(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": [{"model": "gpt-3.5-turbo"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_smuggled_fallback_allowed_under_proxy_wide_opt_in(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [ + {"gpt-4": [{"model": "byok", "api_base": "https://my-byok.example"}]} + ], + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + @pytest.mark.parametrize( + "fallback_field", + ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"], + ) + @pytest.mark.parametrize("surface", ["top_level", "router_settings_override"]) + def test_deeply_nested_api_base_smuggle_rejected_on_both_surfaces(self, fallback_field, surface): + nested = [ + { + "always-fail": [ + { + "model": "x", + fallback_field: [ + {"x": [{"model": "deepseek-chat", "api_base": "http://attacker"}]} + ], + } + ] + } + ] + request_body = {"model": "gpt-4"} + if surface == "top_level": + request_body[fallback_field] = nested + else: + request_body["router_settings_override"] = {fallback_field: nested} + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body=request_body, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_router_settings_override_single_level_api_base_rejected(self): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "router_settings_override": { + "fallbacks": [{"gpt-4": [{"model": "x", "api_base": "http://attacker"}]}] + }, + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_model_less_config_dict_api_base_rejected(self): + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "fallbacks": [{"gpt-4": [{"api_base": "http://attacker"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_nested_api_base_caught_across_router_fallback_rounds(self): + """An ``api_base`` target nested ``ROUTER_MAX_FALLBACKS - 1`` rounds deep + is still reached and rejected.""" + import litellm + + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body=_rounds_deep_api_base_payload(litellm.ROUTER_MAX_FALLBACKS - 1, "fallbacks"), + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_grouping_only_deep_chain_is_rejected_at_depth_limit(self): + """A deep grouping-only chain (``{"g": [{"g": [...]}]}``) is rejected at the + validation-depth limit rather than accepted or raising RecursionError.""" + node: object = ["safe-model"] + for _ in range(5000): + node = [{"grp": node}] + with pytest.raises(ValueError, match="depth"): + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": node}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_pathologically_deep_model_nesting_is_rejected(self): + with pytest.raises(ValueError, match="depth"): + is_request_body_safe( + request_body=_rounds_deep_api_base_payload(5000, "fallbacks"), + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + +class TestIsRequestBodySafeRejectsUrlValuedFallback: + @pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + def test_url_valued_string_fallback_is_rejected(self, fallback_field): + with pytest.raises(ValueError, match="URL-valued fallback"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_field: [{"gpt-4": ["huggingface/http://attacker.example/path"]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + @pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) + def test_url_valued_dict_model_fallback_is_rejected(self, fallback_field): + with pytest.raises(ValueError, match="URL-valued fallback"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + fallback_field: [{"gpt-4": [{"model": "huggingface/http://attacker.example/path"}]}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_ordinary_string_fallback_is_allowed(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": ["gpt-4-backup"]}]}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_ordinary_dict_model_fallback_is_allowed(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": [{"model": "gpt-4-backup"}]}]}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + class TestIsRequestBodySafeBlocksEndpointTargetingFields: """ @@ -1823,6 +2129,46 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride: ) +class TestIsRequestBodySafeBlocksVertexCredentialAlias: + @pytest.mark.parametrize("field", ["vertex_ai_credentials"]) + def test_field_in_request_body_is_rejected(self, field): + with pytest.raises(ValueError, match=field): + is_request_body_safe( + request_body={"model": "gpt-4", field: "attacker-supplied"}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + @pytest.mark.parametrize("field", ["vertex_ai_credentials"]) + def test_admin_opt_in_proxy_wide_allows(self, field): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", field: "byok-supplied"}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_legitimate_request_body_param_still_allowed(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "temperature": 0.7, + "max_tokens": 128, + "user": "end-user-123", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + class TestIsRequestBodySafeBlocksNVCFFunctionOverride: """``nvcf_function_id`` is rejected as a request-body param unless the admin opted in proxy-wide or per-deployment.""" diff --git a/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py index fc0e9aec501..eb1135a240a 100644 --- a/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py +++ b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py @@ -11,12 +11,22 @@ from unittest.mock import AsyncMock, patch import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import iter_request_fallback_targets from litellm.proxy.auth.user_api_key_auth import ( _enforce_key_and_fallback_model_access, - iter_router_fallback_model_names, + _fallback_target_model_name, ) +def _fallback_model_names(fallbacks): + """Model names the auth check validates for a top-level ``fallbacks`` value.""" + return [ + name + for target in iter_request_fallback_targets({"fallbacks": fallbacks}) + if (name := _fallback_target_model_name(target)) is not None + ] + + def _key_with_models(models: List[str]) -> UserAPIKeyAuth: return UserAPIKeyAuth( api_key="hashed", @@ -26,37 +36,40 @@ def _key_with_models(models: List[str]) -> UserAPIKeyAuth: ) -# ── iter_router_fallback_model_names ───────────────────────────────────────── +# ── fallback model-name extraction ─────────────────────────────────────────── -def testiter_router_fallback_model_names_router_config_shape(): +def test_fallback_model_names_router_config_shape(): """Router-config shape: ``[{primary: [fallback_list]}]``.""" - assert list( - iter_router_fallback_model_names( - [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] - ) + assert _fallback_model_names( + [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] ) == ["gpt-4", "claude-3", "o1"] -def testiter_router_fallback_model_names_simple_string_shape(): +def test_fallback_model_names_simple_string_shape(): """Simple top-level shape: list of strings.""" - assert list(iter_router_fallback_model_names(["gpt-4", "claude-3"])) == [ + assert _fallback_model_names(["gpt-4", "claude-3"]) == ["gpt-4", "claude-3"] + + +def test_fallback_model_names_client_side_shape(): + """ClientSideFallbackModel shape: ``[{"model": "..."}]``.""" + assert _fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) == [ "gpt-4", "claude-3", ] -def testiter_router_fallback_model_names_client_side_shape(): - """ClientSideFallbackModel shape: ``[{"model": "..."}]``.""" - assert list( - iter_router_fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) - ) == ["gpt-4", "claude-3"] +def test_fallback_model_names_nested_deployment_fallbacks(): + """A deployment target's own nested fallback field is unrolled too.""" + assert _fallback_model_names( + [{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}] + ) == ["gpt-4", "deepseek-chat"] -def testiter_router_fallback_model_names_empty_or_none(): - assert list(iter_router_fallback_model_names(None)) == [] - assert list(iter_router_fallback_model_names([])) == [] - assert list(iter_router_fallback_model_names("not a list")) == [] +def test_fallback_model_names_empty_or_none(): + assert _fallback_model_names(None) == [] + assert _fallback_model_names([]) == [] + assert _fallback_model_names("not a list") == [] # ── _enforce_key_and_fallback_model_access ──────────────────────────────────── @@ -200,6 +213,98 @@ async def test_top_level_fallback_fields_validated(fallback_field): assert "top-level-smuggled" in seen +@pytest.mark.asyncio +async def test_nested_deployment_fallback_inner_model_validated(): + """A model name nested several fallback rounds deep, inside a deployment + target's own ``fallbacks``, is extracted and passed to can_key_call_model.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "fallbacks": [ + { + "gpt-3.5-turbo": [ + { + "model": "gpt-3.5-turbo", + "fallbacks": [{"gpt-3.5-turbo": ["deep-smuggled-model"]}], + } + ] + } + ], + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert "deep-smuggled-model" in seen + + +@pytest.mark.asyncio +async def test_model_less_fallback_dict_is_skipped_never_passed_as_none(): + """A fallback target dict without a ``model`` key is skipped, never passed + as ``None`` into can_key_call_model / is_valid_fallback_model.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "fallbacks": [ + { + "gpt-3.5-turbo": [ + {"model": "real-fallback"}, + {"api_base": "http://attacker"}, + "string-fallback", + ] + } + ], + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert None not in seen + assert seen == ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + + @pytest.mark.asyncio async def test_router_override_without_fallbacks_does_not_break_auth(): """``router_settings_override`` set without any fallback fields is a diff --git a/tests/test_litellm/proxy/test_provider_url_destination_guard.py b/tests/test_litellm/proxy/test_provider_url_destination_guard.py index 51cd76105d0..c8771abbc8e 100644 --- a/tests/test_litellm/proxy/test_provider_url_destination_guard.py +++ b/tests/test_litellm/proxy/test_provider_url_destination_guard.py @@ -39,6 +39,46 @@ class TestRejectUrlValuedDestinations: assert exc_info.value.status_code == 400 assert exc_info.value.detail["param"] == "model" + def test_provider_prefixed_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "huggingface/https://attacker.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_comma_batch_smuggled_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "gpt-4,huggingface/https://attacker.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_provider_prefixed_uppercase_scheme_url_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"model": "huggingface/HTTPS://evil.example/v1"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_provider_prefixed_plain_model_passes(self): + _reject_url_valued_destinations({"model": "huggingface/BAAI/bge-small-en"}) + + def test_comma_batch_plain_models_pass(self): + _reject_url_valued_destinations({"model": "gpt-4,huggingface/BAAI/bge-small-en"}) + + def test_provider_prefixed_url_respects_allowlist(self, monkeypatch): + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["trusted.example"], + ) + _reject_url_valued_destinations( + {"model": "huggingface/https://trusted.example/v1"} + ) + def test_url_valued_file_id_rejected(self): with pytest.raises(HTTPException) as exc_info: _reject_url_valued_destinations( From 0326722379bb669fbb910df9132e2b760d9b15d8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 18:04:25 -0700 Subject: [PATCH 023/130] docs(issue-template): ask for a numbered list of reproduction steps (#34207) --- .github/ISSUE_TEMPLATE/bug_report.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index bbe4b76775d..665f8456f0b 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -30,7 +30,7 @@ body: id: steps-to-reproduce attributes: label: Steps to Reproduce - description: Please provide detailed steps to reproduce this bug(A curl/python code to reproduce the bug) + description: Please provide a numbered list of the exact steps to reproduce this bug (include a curl/python snippet to reproduce it). Number each step (1., 2., 3., ...) in the order you performed them. placeholder: | 1. config.yaml file/ .env file/ etc. 2. Run the following code... From 5b676b91bde093f58a7301efff3d4f3f4da1003e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 18:14:24 -0700 Subject: [PATCH 024/130] test(ui): fix key and credential e2e specs after the overflow menu migrations (#34206) Delete Key moved into the key info page's overflow dropdown (#34116) and the credentials table's row actions moved into a shared DataTable overflow menu, so both specs were clicking a button that no longer exists. Point them at the menu items instead. Add a CredentialsPanel unit test asserting the update payload drops the masked api key and keeps the edited api base, so that guard is not held up solely by an e2e a table migration can silently disarm. --- .../tests/modelsPage/credentials.spec.ts | 3 +- .../e2e_tests/tests/proxy-admin/keys.spec.ts | 3 +- .../model_add/CredentialsPanel.test.tsx | 36 ++++++++++++++++--- 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts index 8b7824813a4..7c836068567 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts @@ -38,7 +38,8 @@ test.describe("Edit LLM credential", () => { const row = page.locator("tr", { hasText: credentialName }); await expect(row).toBeVisible({ timeout: 15_000 }); - await row.getByRole("button").first().click(); + await row.getByTestId(`credential-actions-${credentialName}`).click(); + await page.getByTestId("credential-action-edit").click(); const modal = page.locator(".ant-modal-content").filter({ hasText: "Edit Credential" }); await expect(modal).toBeVisible({ timeout: 10_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts index a55c19a53de..c44957ea737 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts @@ -103,7 +103,8 @@ test.describe("Proxy Admin - Keys", () => { await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); - await page.getByRole("button", { name: "Delete Key" }).click(); + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Delete Key" }).click(); const modal = page.locator(".ant-modal:visible"); await expect(modal).toBeVisible({ timeout: 5_000 }); diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialsPanel.test.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialsPanel.test.tsx index af38645b00d..93a015c8a36 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialsPanel.test.tsx @@ -4,7 +4,7 @@ import userEvent from "@testing-library/user-event"; import { UploadProps } from "antd/es/upload"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { CredentialItem, credentialCreateCall } from "@/components/networking"; +import { CredentialItem, credentialCreateCall, credentialUpdateCall } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; import CredentialsPanel from "./CredentialsPanel"; @@ -51,11 +51,17 @@ vi.mock("./CredentialModal", () => ({ if (!open) { return null; } + const values = + mode === "edit" + ? { + credential_name: "openai-key", + custom_llm_provider: "openai", + api_key: "sk-1****2345", + api_base: "https://proxy.e2e.example.com/v1", + } + : { credential_name: "new-cred", custom_llm_provider: "openai" }; return ( - ); @@ -179,6 +185,26 @@ describe("CredentialsPanel", () => { expect(NotificationsManager.success).not.toHaveBeenCalled(); }); + it("drops the masked api key from the update payload while keeping the edited api base", async () => { + const user = userEvent.setup(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + mockUseCredentials.mockReturnValue({ data: { credentials }, isLoading: false, refetch: vi.fn() }); + vi.mocked(credentialUpdateCall).mockResolvedValueOnce(undefined as never); + + renderPanel(); + + await user.click(screen.getByTestId("credential-actions-openai-key")); + await user.click(await screen.findByTestId("credential-action-edit")); + await user.click(screen.getByTestId("credential-modal-edit-submit")); + + await waitFor(() => { + expect(credentialUpdateCall).toHaveBeenCalled(); + }); + const [, updatedName, payload] = vi.mocked(credentialUpdateCall).mock.calls[0]; + expect(updatedName).toBe("openai-key"); + expect(payload.credential_values).toEqual({ api_base: "https://proxy.e2e.example.com/v1" }); + }); + describe("Admin Viewer write-action gating", () => { // Admin Viewer can VIEW credentials but must not add / edit / delete them. it("hides the Add Credential button but still lists credentials", () => { From 82d3116be90f290dc7b24fa0657021a8fea599b2 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 21 Jul 2026 18:36:26 -0700 Subject: [PATCH 025/130] fix(ui): reflect REDIS_* env cache config and stop the UI overwriting the stored password (#34160) The Cache Settings page read only the database row, so a response cache pointed at Redis purely through REDIS_* env vars showed a blank page while the cache worked. It also masked credentials on read with a partial-reveal string and re-persisted whatever the form submitted, so an admin who edited an unrelated field and pressed Save wrote the mask string over the real Redis password, breaking auth. GET /cache/settings now overlays the same REDIS_* kwargs the runtime resolves from when the stored config leaves a field unset, and redacts credentials with a fixed marker. POST /cache/settings restores the stored secret behind any credential echoed back as the marker or omitted, and drops an env-sourced marker rather than persisting it; the response no longer echoes plaintext credentials. The connection test resolves a redacted credential back to the stored value the same way. The dashboard never prefills a credential and drops the marker from the save payload, mirroring the Coordination Redis tab. Resolves LIT-4315 --- .../cache_settings_endpoints.py | 257 +++++++-- .../test_cache_settings_endpoints.py | 512 ++++++++++++++++++ .../cache_settings/CacheFieldSection.tsx | 9 +- .../cache_settings/CacheFormField.tsx | 21 +- .../cache_settings/cacheSettingsFields.ts | 10 + .../cache_settings/cacheSettingsUtils.test.ts | 48 +- .../cache_settings/cacheSettingsUtils.ts | 20 +- .../_components/cache_settings/index.tsx | 6 +- 8 files changed, 838 insertions(+), 45 deletions(-) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 9f45cb619aa..7c0d8958a28 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -18,8 +18,9 @@ from pydantic import BaseModel, Field import litellm from litellm._logging import verbose_proxy_logger +from litellm._redis import _redis_kwargs_from_environment from litellm._uuid import uuid -from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ( AUDIT_ACTIONS, LiteLLM_AuditLogs, @@ -43,6 +44,17 @@ router = APIRouter() # (e.g. redis://:secret@host:6379/1). _CACHE_SENSITIVE_FIELDS: set = {"password", "sentinel_password", "url"} +# The env fallback resolves the full set of redis.Redis kwargs, which includes +# credential-bearing params (azure_client_secret, ssl_password, ...) that are +# not cache UI fields. Only overlay fields the settings page actually renders, +# so the read never surfaces a credential the UI does not manage. +_CACHE_SETTINGS_FIELD_NAMES: frozenset = frozenset(field.field_name for field in CACHE_SETTINGS_FIELDS) + +# Classifier used, alongside _CACHE_SENSITIVE_FIELDS, to redact any +# credential-bearing key before it leaves the server (`url` is kept in the +# explicit set because its name carries no sensitive segment). +_CREDENTIAL_CLASSIFIER = SensitiveDataMasker() + _REDACTED_VALUE = "***REDACTED***" @@ -67,6 +79,165 @@ def _resolve_cache_url_precedence(settings: Mapping[str, Any]) -> dict[str, Any] return {k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS} +def _parse_stored_settings(cache_settings_value: object) -> dict[str, Any]: + """Normalize a stored cache_settings blob to a dict. + + The prisma column comes back as either a JSON string or an already-parsed + dict depending on the client, so callers that json.loads unconditionally + silently drop the whole (still-encrypted) row on the dict path. + """ + parsed = json.loads(cache_settings_value) if isinstance(cache_settings_value, str) else cache_settings_value + return parsed if isinstance(parsed, dict) else {} + + +def _overlay_environment(stored: Mapping[str, Any]) -> dict[str, Any]: + """Fill connection fields from the REDIS_* environment the cache actually reads. + + A response cache pointed at Redis resolves host/port/password/etc. from the + REDIS_* env vars when the stored config leaves them unset, so a cache + configured purely through the environment works while its settings page, + which reads only the database row, shows blank. Overlaying the same env + kwargs the runtime uses makes the page reflect the effective connection. + Stored values win; the environment only fills what the stored config omits. + """ + env_kwargs = { + key: value for key, value in _redis_kwargs_from_environment().items() if key in _CACHE_SETTINGS_FIELD_NAMES + } + if not env_kwargs: + return dict(stored) + effective = {**env_kwargs, **stored} + # the env fallback is a Redis connection, so name the type when the stored + # config did not, letting the UI render the Redis fields it just populated + effective.setdefault("type", "redis") + return effective + + +def _redact_credentials(settings: Mapping[str, Any]) -> dict[str, Any]: + """Replace credential-bearing values with a fixed marker, keeping the rest. + + The marker is unambiguous on the way back in: an admin who edits an + unrelated field and re-submits sends the marker for the untouched secret, + which the update path maps back to the stored value rather than persisting + the marker over a working password. + """ + return { + key: (_REDACTED_VALUE if value is not None and _is_credential_field(key) else value) + for key, value in settings.items() + } + + +def _is_credential_field(key: str) -> bool: + """Whether a cache setting carries a credential and must be redacted on read.""" + return key in _CACHE_SENSITIVE_FIELDS or _CREDENTIAL_CLASSIFIER.is_sensitive_key(key) + + +def _has_connection_target(value: object) -> bool: + """Whether a payload value names a live discrete connection target.""" + if isinstance(value, str): + return value.strip() != "" and value != _REDACTED_VALUE + return value not in (None, [], {}) + + +# Every field that identifies which Redis a credential belongs to, across node +# (host/port/url), cluster (redis_startup_nodes), and sentinel +# (sentinel_nodes/service_name) modes. A stored secret is bound to these. +_CONNECTION_TARGET_FIELDS: tuple = ( + "host", + "port", + "url", + "redis_startup_nodes", + "sentinel_nodes", + "service_name", +) + + +def _target_repr(value: object) -> str: + """Canonical string form of a connection-target value for equality checks. + + The client may serialize the same target differently from storage (a port as + "6379" vs 6379, node lists round-tripped through JSON), so compare normalized + forms rather than raw values to avoid treating an unchanged target as a change. + """ + if isinstance(value, (list, dict)): + return json.dumps(value, sort_keys=True, default=str) + return str(value) + + +def _saved_secret_is_reusable(incoming: Mapping[str, Any], saved: Mapping[str, Any]) -> bool: + """Whether a stored credential may be restored for this request. + + A stored secret belongs to the stored connection target, so it is reused only + when the request describes that same target on every dimension the stored + config pins (host/port, url, cluster nodes, sentinel nodes/service). This + prevents credential replay: a caller cannot omit the credential, point at a + different (or incomplete) target, and have the proxy send the stored secret + to a Redis of their choosing. + + Non-secret target fields (host/port/nodes/service) must be supplied and match + in normalized form, so equivalent representations (port "6379" vs 6379) are + not seen as a change while an omitted or different value is. ``url`` is the + exception: it is itself the secret and the form never re-prefills it, so a + redacted or omitted url means "keep the stored url" (same target) and only a + different supplied url blocks reuse. + """ + for field in _CONNECTION_TARGET_FIELDS: + saved_value = saved.get(field) + if saved_value in (None, "", [], {}): + continue # the stored config does not pin this dimension + incoming_value = incoming.get(field) + if field == "url": + if incoming_value in (None, "", _REDACTED_VALUE): + continue # url kept as-is (same target) + if _target_repr(incoming_value) != _target_repr(saved_value): + return False + continue + if _target_repr(incoming_value) != _target_repr(saved_value): + return False # a pinned target field is missing or different + return True + + +def _merge_over_saved(incoming: Mapping[str, Any], saved: Mapping[str, Any]) -> dict[str, Any]: + """Keep the stored secret behind any credential the caller echoed back redacted or omitted. + + GET returns credentials as the marker and the form never re-prefills a + secret, so a save that does not touch a credential arrives with the marker + or with the field absent. Either way the real secret must survive: it is + restored from the stored row, or dropped when there is no stored row (the + value is env-sourced and the marker must never be persisted). Non-secret + fields are taken from the incoming payload as-is, so clearing one still works. + + ``url`` is the exception: it is credential-bearing (redacted) yet also a + connection-mode selector that url-precedence resolves against host/port. If + the caller supplies a discrete target (host, cluster, or sentinel nodes), a + stored url is a stale mode the caller is leaving, so it is dropped rather + than restored, otherwise url-precedence would resurrect it and discard the + submitted host/port. + """ + switching_to_discrete_target = ( + _has_connection_target(incoming.get("host")) + or _has_connection_target(incoming.get("redis_startup_nodes")) + or _has_connection_target(incoming.get("sentinel_nodes")) + ) + reuse_saved_secret = _saved_secret_is_reusable(incoming, saved) + merged = dict(incoming) + for field in _CACHE_SENSITIVE_FIELDS: + # A value the caller explicitly supplied is honored verbatim: a new + # secret, or an empty string / null to clear the stored one. Only an + # omitted field or the echoed-back marker triggers preserve-or-drop. + if field in incoming and incoming[field] != _REDACTED_VALUE: + continue + if field == "url" and switching_to_discrete_target: + merged.pop(field, None) + continue + if field in saved and reuse_saved_secret: + merged[field] = saved[field] + else: + # nothing stored to reuse, or the caller is pointing at a different + # target: never persist/replay the marker or the stored secret + merged.pop(field, None) + return merged + + def _redact_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]: """Replace every value in a settings map with a fixed marker. @@ -270,34 +441,34 @@ async def get_cache_settings( # Get cache settings fields from types file cache_fields = [field.model_copy(deep=True) for field in CACHE_SETTINGS_FIELDS] - # Try to get cache settings from database - current_values = {} + # Read the stored settings (decrypted); an env-only cache has none. + stored: dict[str, Any] = {} if prisma_client is not None: cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) if cache_config is not None and cache_config.cache_settings: - # Decrypt cache settings - cache_settings_json = cache_config.cache_settings - if isinstance(cache_settings_json, str): - cache_settings_dict = json.loads(cache_settings_json) - else: - cache_settings_dict = cache_settings_json + stored = proxy_config._decrypt_db_variables( + variables_dict=_parse_stored_settings(cache_config.cache_settings) + ) - # Decrypt environment variables - decrypted_settings = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict) + # Fill connection fields from the REDIS_* environment the cache resolves + # from when the stored config leaves them unset, then apply url precedence + # so a url-mode config does not surface conflicting discrete fields (which + # would otherwise let a no-op save silently switch it to host/port). + effective = _resolve_cache_url_precedence(_overlay_environment(stored)) - # Derive redis_type for UI based on settings - # UI uses redis_type to show/hide fields, backend only stores 'type' - if decrypted_settings.get("type") == "redis": - if decrypted_settings.get("redis_startup_nodes"): - decrypted_settings["redis_type"] = "cluster" - elif decrypted_settings.get("sentinel_nodes"): - decrypted_settings["redis_type"] = "sentinel" - else: - decrypted_settings["redis_type"] = "node" + # Derive redis_type for UI based on settings + # UI uses redis_type to show/hide fields, backend only stores 'type' + if effective.get("type") == "redis": + if effective.get("redis_startup_nodes"): + effective["redis_type"] = "cluster" + elif effective.get("sentinel_nodes"): + effective["redis_type"] = "sentinel" + else: + effective["redis_type"] = "node" - # Mask credential fields so the GET response never carries - # plaintext Redis / Sentinel passwords off the server. - current_values = mask_sensitive_keys(decrypted_settings, _CACHE_SENSITIVE_FIELDS) + # Redact credential fields so the GET response never carries a plaintext + # Redis / Sentinel password off the server. + current_values = _redact_credentials(effective) # Update field values with current values for field in cache_fields: @@ -331,10 +502,27 @@ async def test_cache_connection( to verify the credentials work without affecting global state. """ from litellm import Cache + from litellm.proxy.proxy_server import prisma_client, proxy_config try: - cache_settings = _resolve_cache_url_precedence(request.cache_settings) - verbose_proxy_logger.debug("Testing cache connection with settings: %s", cache_settings) + # A credential the form left untouched arrives redacted; resolve it back + # to the stored secret so the test connects with the real password. A + # lookup failure must not block the test, so fall back to no stored row. + saved_settings: dict[str, Any] = {} + if prisma_client is not None: + try: + existing_row = await CacheConfigRepository(prisma_client).table.find_unique( + where={"id": "cache_config"} + ) + if existing_row is not None and existing_row.cache_settings: + saved_settings = proxy_config._decrypt_db_variables( + variables_dict=_parse_stored_settings(existing_row.cache_settings) + ) + except Exception: # noqa: BLE001 - a saved-settings lookup failure must not block a connection test + saved_settings = {} + cache_settings = _resolve_cache_url_precedence(_merge_over_saved(request.cache_settings, saved_settings)) + # cache_settings now carries the resolved plaintext credential; never log it raw + verbose_proxy_logger.debug("Testing cache connection with settings: %s", _redact_credentials(cache_settings)) # Only support Redis for now if cache_settings.get("type") != "redis": @@ -400,19 +588,20 @@ async def update_cache_settings( ) try: - cache_settings = _resolve_cache_url_precedence(request.cache_settings) - - # Snapshot the prior settings (key set only — values get redacted in - # the audit row) so the audit-log entry shows which fields changed. + # Read the stored row first: its decrypted values back any credential the + # caller echoed back redacted, and its key set drives the audit diff. existing_row = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) before_settings: Optional[Dict[str, Any]] = None + saved_settings: dict[str, Any] = {} if existing_row is not None and existing_row.cache_settings: - try: - before_settings = json.loads(existing_row.cache_settings) - except (TypeError, ValueError): - before_settings = None + before_settings = _parse_stored_settings(existing_row.cache_settings) + saved_settings = proxy_config._decrypt_db_variables(variables_dict=before_settings) action: AUDIT_ACTIONS = "updated" if existing_row is not None else "created" + # Preserve stored secrets behind any redacted or omitted credential, then + # resolve the url-vs-discrete-fields precedence. + cache_settings = _resolve_cache_url_precedence(_merge_over_saved(request.cache_settings, saved_settings)) + # Encrypt sensitive fields (keep redis_type for storage) encrypted_settings = proxy_config._encrypt_env_variables(environment_variables=cache_settings) @@ -461,7 +650,7 @@ async def update_cache_settings( return { "message": "Cache settings updated successfully", "status": "success", - "settings": cache_settings, + "settings": _redact_credentials(cache_settings), } except Exception as e: verbose_proxy_logger.error(f"Error updating cache settings: {str(e)}") diff --git a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py index f4c6d4f8d15..2504b5744fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py @@ -17,9 +17,14 @@ from litellm.proxy._types import LitellmTableNames, LitellmUserRoles from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.management_endpoints.cache_settings_endpoints import ( _CACHE_SENSITIVE_FIELDS, + _REDACTED_VALUE, CacheSettingsManager, CacheSettingsUpdateRequest, CacheTestRequest, + _merge_over_saved, + _overlay_environment, + _parse_stored_settings, + _redact_credentials, _resolve_cache_url_precedence, get_cache_settings, test_cache_connection, @@ -610,3 +615,510 @@ async def test_update_cache_settings_no_audit_when_disabled(monkeypatch): ) assert audit_calls == [] + + +class TestParseStoredSettings: + """The stored blob arrives as a JSON string or a parsed dict; both must + normalize to a dict so the secret-preservation read never silently drops it.""" + + def test_parses_a_json_string(self): + assert _parse_stored_settings('{"host": "h", "password": "pw"}') == {"host": "h", "password": "pw"} + + def test_passes_a_dict_through(self): + assert _parse_stored_settings({"host": "h", "password": "pw"}) == {"host": "h", "password": "pw"} + + def test_non_mapping_becomes_empty(self): + assert _parse_stored_settings(None) == {} + assert _parse_stored_settings("[1, 2]") == {} + + +class TestMergeOverSaved: + """The secret-preservation contract behind the redacted-resubmit fix.""" + + def test_redacted_secret_restores_stored_value(self): + # same connection target, an unrelated field edited: the stored secret + # is restored behind the redacted resubmit + merged = _merge_over_saved( + incoming={"type": "redis", "host": "samehost", "namespace": "new", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "samehost", "password": "realpw"}, + ) + assert merged["namespace"] == "new" + assert merged["password"] == "realpw" + + def test_stored_secret_not_replayed_to_a_different_target(self): + # credential replay guard: omitting the password while pointing at a new + # host must NOT resurrect the stored secret (it would be sent elsewhere) + merged = _merge_over_saved( + incoming={"type": "redis", "host": "attacker.example.com", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "real-redis", "password": "realpw"}, + ) + assert "password" not in merged + + def test_omitted_secret_restores_stored_value(self): + # same host (target unchanged), password field omitted entirely + merged = _merge_over_saved( + incoming={"type": "redis", "host": "samehost", "namespace": "n"}, + saved={"type": "redis", "host": "samehost", "password": "realpw"}, + ) + assert merged["password"] == "realpw" + + def test_sentinel_password_not_replayed_to_different_sentinel_nodes(self): + # sentinel target change with an omitted sentinel_password must not + # resurrect the stored one and send it to the caller's sentinels + merged = _merge_over_saved( + incoming={"type": "redis", "sentinel_nodes": [["attacker", 26379]], "service_name": "mymaster"}, + saved={ + "type": "redis", + "sentinel_nodes": [["real", 26379]], + "service_name": "mymaster", + "sentinel_password": "realsp", + }, + ) + assert "sentinel_password" not in merged + + def test_sentinel_password_preserved_when_sentinel_target_unchanged(self): + merged = _merge_over_saved( + incoming={"type": "redis", "sentinel_nodes": [["real", 26379]], "service_name": "mymaster"}, + saved={ + "type": "redis", + "sentinel_nodes": [["real", 26379]], + "service_name": "mymaster", + "sentinel_password": "realsp", + }, + ) + assert merged["sentinel_password"] == "realsp" + + def test_password_not_replayed_to_different_cluster_nodes(self): + merged = _merge_over_saved( + incoming={"type": "redis", "redis_startup_nodes": [{"host": "attacker", "port": "7001"}]}, + saved={ + "type": "redis", + "redis_startup_nodes": [{"host": "real", "port": "7001"}], + "password": "realpw", + }, + ) + assert "password" not in merged + + def test_equivalent_target_representations_still_preserve_secret(self): + # the client sends port as a string, storage holds it as an int: the + # target is unchanged, so the untouched password must not be dropped + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "port": "6379", "password": _REDACTED_VALUE}, + saved={"type": "redis", "host": "h", "port": 6379, "password": "realpw"}, + ) + assert merged["password"] == "realpw" + + def test_explicit_empty_string_clears_the_secret(self): + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": ""}, + saved={"type": "redis", "host": "h", "password": "realpw"}, + ) + assert merged.get("password") == "" + + def test_explicit_null_clears_the_secret(self): + # an explicit null is a clear, not an omission, so it must not restore + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": None}, + saved={"type": "redis", "host": "h", "password": "realpw"}, + ) + assert merged.get("password") is None + + def test_secret_not_reused_when_a_pinned_target_field_is_omitted(self): + # omitting the host (a pinned target) means the request does not describe + # the stored target, so the stored secret must not be restored (and thus + # cannot be sent to whatever host the incomplete request resolves to) + merged = _merge_over_saved( + incoming={"type": "redis", "port": "6379"}, + saved={"type": "redis", "host": "real", "port": 6379, "password": "realpw"}, + ) + assert "password" not in merged + + def test_redacted_secret_with_no_stored_value_is_dropped(self): + # env-sourced secret: nothing stored to restore, so the marker must not + # be persisted; the environment stays the source at runtime + merged = _merge_over_saved( + incoming={"type": "redis", "host": "h", "password": _REDACTED_VALUE}, + saved={}, + ) + assert "password" not in merged + + def test_new_secret_value_wins(self): + merged = _merge_over_saved( + incoming={"password": "brandnewpw"}, + saved={"password": "realpw"}, + ) + assert merged["password"] == "brandnewpw" + + def test_switching_from_url_to_host_port_drops_stored_url(self): + # admin migrates a url-mode cache to discrete host/port: the stored url + # must not be resurrected (url precedence would then discard host/port) + merged = _merge_over_saved( + incoming={"type": "redis", "host": "newhost", "port": "6379"}, + saved={"type": "redis", "url": "redis://:pw@oldhost:6379/0"}, + ) + assert "url" not in merged + assert merged["host"] == "newhost" + assert merged["port"] == "6379" + + def test_untouched_url_is_preserved_without_a_discrete_target(self): + # a url-mode save that touches nothing keeps the stored url + merged = _merge_over_saved( + incoming={"type": "redis", "namespace": "ns"}, + saved={"type": "redis", "url": "redis://:pw@host:6379/0"}, + ) + assert merged["url"] == "redis://:pw@host:6379/0" + + +def test_overlay_environment_fills_unset_connection_fields(monkeypatch): + """A cache with no stored connection resolves REDIS_* env for the UI.""" + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_PORT", "6380") + monkeypatch.setenv("REDIS_PASSWORD", "env-password") + + effective = _overlay_environment({}) + + assert effective["host"] == "redis.internal" + assert effective["port"] == "6380" + assert effective["password"] == "env-password" + assert effective["type"] == "redis" + + +def test_overlay_environment_stored_value_wins(monkeypatch): + monkeypatch.setenv("REDIS_HOST", "env-host") + effective = _overlay_environment({"type": "redis", "host": "stored-host"}) + assert effective["host"] == "stored-host" + + +@pytest.mark.asyncio +async def test_get_cache_settings_falls_back_to_redis_env(monkeypatch): + """A cache configured purely through REDIS_* env vars shows its effective + connection instead of a blank page, with the password redacted.""" + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_PORT", "6380") + monkeypatch.setenv("REDIS_PASSWORD", "env-password") + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values["host"] == "redis.internal" + assert values["port"] == "6380" + assert values["type"] == "redis" + # the env password is surfaced as configured, not leaked in plaintext + assert values["password"] == _REDACTED_VALUE + + +@pytest.mark.asyncio +async def test_get_cache_settings_redacts_password_with_marker(monkeypatch): + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + cache_row = MagicMock() + cache_row.cache_settings = json.dumps( + {"type": "redis", "host": "h", "password": "supersecret", "namespace": "ns"} + ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=cache_row) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + assert response.current_values["password"] == _REDACTED_VALUE + assert response.current_values["namespace"] == "ns" + + +@pytest.mark.asyncio +async def test_get_cache_settings_url_mode_hides_env_discrete_fields(monkeypatch): + """A url-mode stored config must not surface env-overlaid host/port. + + Otherwise a no-op save would submit the env host and, via url precedence, + silently switch the cache off its configured url. + """ + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_USERNAME"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "env-host") + monkeypatch.setenv("REDIS_PORT", "6380") + + cache_row = MagicMock() + cache_row.cache_settings = {"type": "redis", "url": "redis://:pw@stored-host:6379/0"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=cache_row) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values["url"] == _REDACTED_VALUE + # the env host/port must not leak in and shadow the url + assert "host" not in values + assert "port" not in values + + +def _mock_proxy_config_identity_crypto(): + proxy_config = MagicMock() + proxy_config._encrypt_env_variables = MagicMock( + side_effect=lambda environment_variables: dict(environment_variables) + ) + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + proxy_config._init_cache = MagicMock() + proxy_config.switch_on_llm_response_caching = MagicMock() + return proxy_config + + +@pytest.mark.asyncio +async def test_update_preserves_stored_password_on_redacted_resubmit(monkeypatch): + """Editing an unrelated field and re-submitting the redacted password must + keep the stored secret, not persist the marker over a working password.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + existing = MagicMock() + # prisma returns the Json column as an already-parsed dict, not a JSON + # string; a reader that json.loads unconditionally would drop the whole row + existing.cache_settings = {"type": "redis", "host": "oldhost", "password": "realpw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + result = await update_cache_settings( + request=CacheSettingsUpdateRequest( + # same host (the target is unchanged), an unrelated field edited + cache_settings={"type": "redis", "host": "oldhost", "namespace": "edited", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert persisted["host"] == "oldhost" + assert persisted["namespace"] == "edited" + assert persisted["password"] == "realpw" + # the response never echoes the plaintext secret back either + assert result["settings"]["password"] == _REDACTED_VALUE + + +@pytest.mark.asyncio +async def test_update_drops_env_sourced_redacted_secret(monkeypatch): + """With no stored row, a re-submitted redacted secret is env-sourced; the + marker must not be persisted so the environment stays the source.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await update_cache_settings( + request=CacheSettingsUpdateRequest( + cache_settings={"type": "redis", "host": "h", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert "password" not in persisted + + +@pytest.mark.asyncio +async def test_update_applies_new_password(monkeypatch): + """A real new secret value replaces the stored one.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + + existing = MagicMock() + existing.cache_settings = json.dumps({"type": "redis", "host": "h", "password": "oldpw"}) + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + mock_prisma.db.litellm_cacheconfig.upsert = AsyncMock() + proxy_config = _mock_proxy_config_identity_crypto() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await update_cache_settings( + request=CacheSettingsUpdateRequest( + cache_settings={"type": "redis", "host": "h", "password": "brandnewpw"} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = proxy_config._encrypt_env_variables.call_args.kwargs["environment_variables"] + assert persisted["password"] == "brandnewpw" + + +@pytest.mark.asyncio +async def test_test_cache_connection_survives_saved_lookup_failure(monkeypatch): + """A failed saved-settings lookup must not block the connection test. + + The test endpoint reads the stored row to resolve a redacted credential, but + that read can raise (a misconfigured or unavailable client), and it must fall + back to the submitted settings rather than abort — otherwise a shared client + left in an odd state by another test would break every connection test. + """ + monkeypatch.setattr(litellm, "store_audit_logs", False) + + # a client whose find_unique is not awaitable, so the saved read raises + bad_prisma = MagicMock() + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", bad_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + ): + mock_cache_class.return_value = cache_instance + result = await test_cache_connection( + request=CacheTestRequest(cache_settings={"type": "redis", "host": "h", "port": "6379", "password": "pw"}), + user_api_key_dict=_admin_auth(), + ) + + mock_cache_class.assert_called_once() + assert result.status == "success" + + +@pytest.mark.asyncio +async def test_get_cache_settings_does_not_surface_non_display_env_credentials(monkeypatch): + """The env overlay must not leak credential kwargs the UI does not manage. + + _redis_kwargs_from_environment resolves every redis.Redis kwarg, including + secrets like azure_client_secret; only cache display fields may be surfaced, + so a non-admin reading /cache/settings never retrieves such a credential. + """ + for var in ("REDIS_URL", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_AZURE_CLIENT_SECRET"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("REDIS_HOST", "redis.internal") + monkeypatch.setenv("REDIS_AZURE_CLIENT_SECRET", "super-azure-secret") + + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=None) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + ): + response = await get_cache_settings(user_api_key_dict=_admin_auth()) + + values = response.current_values + assert values.get("host") == "redis.internal" + # the non-display credential must not appear in the response at all + assert "azure_client_secret" not in values + assert "super-azure-secret" not in values.values() + + +@pytest.mark.asyncio +async def test_test_cache_connection_does_not_log_plaintext_credentials(monkeypatch, caplog): + """The connection test must not write the resolved plaintext secret to logs. + + _merge_over_saved substitutes the stored password for a redacted resubmit, so + the settings dict carries the real secret; the debug log must redact it. + """ + import logging + + existing = MagicMock() + existing.cache_settings = {"type": "redis", "host": "h", "port": "6379", "password": "realredispw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"), + ): + mock_cache_class.return_value = cache_instance + # resubmit the redacted marker; the merge resolves it to the stored secret + await test_cache_connection( + request=CacheTestRequest( + cache_settings={"type": "redis", "host": "h", "port": "6379", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + ) + + # the real password was used to build the client but never written to the log + assert mock_cache_class.call_args.kwargs["password"] == "realredispw" + assert "realredispw" not in caplog.text + + +@pytest.mark.asyncio +async def test_test_cache_connection_does_not_replay_saved_password_to_new_host(monkeypatch): + """Credential-replay guard on the connection test. + + A caller that submits a different host while omitting the password must not + have the stored password restored and sent to the caller-chosen host. + """ + existing = MagicMock() + existing.cache_settings = {"type": "redis", "host": "real-redis", "port": "6379", "password": "realredispw"} + mock_prisma = MagicMock() + mock_prisma.db.litellm_cacheconfig.find_unique = AsyncMock(return_value=existing) + proxy_config = MagicMock() + proxy_config._decrypt_db_variables = MagicMock(side_effect=lambda variables_dict: dict(variables_dict)) + + cache_instance = MagicMock() + cache_instance.cache = MagicMock() + cache_instance.cache.test_connection = AsyncMock(return_value={"status": "success", "message": "ok"}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", proxy_config), + patch("litellm.Cache") as mock_cache_class, + ): + mock_cache_class.return_value = cache_instance + await test_cache_connection( + request=CacheTestRequest( + cache_settings={"type": "redis", "host": "attacker.example.com", "port": "6379"} + ), + user_api_key_dict=_admin_auth(), + ) + + called_kwargs = mock_cache_class.call_args.kwargs + # the stored password is NOT sent to the attacker-chosen host + assert called_kwargs.get("password") != "realredispw" + assert "password" not in called_kwargs diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFieldSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFieldSection.tsx index ced822cd796..96106869009 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFieldSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFieldSection.tsx @@ -10,6 +10,7 @@ interface CacheFieldSectionProps { embeddingModels: EmbeddingModelOption[]; gridCols?: string; headingLevel?: "h4" | "h5"; + configuredSecrets?: ReadonlySet; } const CacheFieldSection: React.FC = ({ @@ -19,6 +20,7 @@ const CacheFieldSection: React.FC = ({ embeddingModels, gridCols = "grid-cols-1 gap-6 sm:grid-cols-2", headingLevel = "h4", + configuredSecrets, }) => { const fields = fieldsForSection(section, redisType); if (fields.length === 0) { @@ -32,7 +34,12 @@ const CacheFieldSection: React.FC = ({ {title}
{fields.map((field) => ( - + ))}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx index d92ca302901..dbd8c32d18d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx @@ -7,22 +7,29 @@ export interface EmbeddingModelOption { label: string; } +export const SECRET_ALREADY_SET_PLACEHOLDER = "Already set. Enter a new value to replace it."; + interface CacheFormFieldProps { field: CacheField; embeddingModels: EmbeddingModelOption[]; + isSecretConfigured?: boolean; } -const renderControl = (field: CacheField, embeddingModels: EmbeddingModelOption[]): React.ReactNode => { +const renderControl = ( + field: CacheField, + embeddingModels: EmbeddingModelOption[], + placeholder: string, +): React.ReactNode => { switch (field.type) { case "boolean": return ; case "password": - return ; + return ; case "integer": case "float": - return ; + return ; case "list": - return ; + return ; case "model-select": return ( ; + return ; } }; -const CacheFormField: React.FC = ({ field, embeddingModels }) => ( +const CacheFormField: React.FC = ({ field, embeddingModels, isSecretConfigured = false }) => ( = ({ field, embeddingModels rules={field.rules} valuePropName={field.type === "boolean" ? "checked" : "value"} > - {renderControl(field, embeddingModels)} + {renderControl(field, embeddingModels, isSecretConfigured ? SECRET_ALREADY_SET_PLACEHOLDER : field.helpText)} ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts index 1f5b566fc5f..e33b525c3ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts @@ -8,6 +8,10 @@ export type CacheSection = "connection" | "cluster" | "sentinel" | "semantic" | export type CacheFieldRule = NonNullable[number]; +// Marker the backend returns for a configured credential and maps back to the +// stored secret on save, so the plaintext never round-trips through the form. +export const REDACTED_VALUE = "***REDACTED***"; + export interface CacheField { readonly name: string; readonly label: string; @@ -17,6 +21,9 @@ export interface CacheField { readonly redisType: RedisType | null; readonly defaultValue?: string | number | boolean; readonly rules?: CacheFieldRule[]; + // Credential field: never prefilled into the form, and dropped from the save + // payload when left untouched so the redacted marker is never persisted. + readonly secret?: boolean; } export const REDIS_TYPES: readonly RedisType[] = ["node", "cluster", "sentinel", "semantic"]; @@ -93,6 +100,7 @@ export const CACHE_FIELDS: readonly CacheField[] = [ helpText: "Full Redis/Valkey connection URL (e.g. redis://:password@host:6379/1). When set, it takes precedence over Host, Port, Password, and Database Index.", redisType: null, + secret: true, }, { name: "host", @@ -128,6 +136,7 @@ export const CACHE_FIELDS: readonly CacheField[] = [ section: "connection", helpText: "Redis server password", redisType: null, + secret: true, }, { name: "username", @@ -170,6 +179,7 @@ export const CACHE_FIELDS: readonly CacheField[] = [ section: "sentinel", helpText: "Password for Redis Sentinel authentication", redisType: "sentinel", + secret: true, }, { name: "similarity_threshold", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts index 79f28a97842..c530519ee06 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts @@ -1,5 +1,6 @@ import { describe, it, expect } from "vitest"; -import { buildCachePayload, buildInitialValues, fieldsForSection } from "./cacheSettingsUtils"; +import { buildCachePayload, buildInitialValues, configuredSecretFields, fieldsForSection } from "./cacheSettingsUtils"; +import { REDACTED_VALUE } from "./cacheSettingsFields"; describe("fieldsForSection", () => { it("should only include a redis-type-specific field when that type is selected", () => { @@ -83,4 +84,49 @@ describe("buildCachePayload", () => { const payload = buildCachePayload("node", { sentinel_nodes: '[["localhost",26379]]' }, { forTesting: false }); expect(payload).not.toHaveProperty("sentinel_nodes"); }); + + it("should drop a secret whose value is the redacted marker so it is never persisted", () => { + const payload = buildCachePayload( + "node", + { host: "localhost", password: REDACTED_VALUE, url: REDACTED_VALUE }, + { forTesting: false }, + ); + expect(payload).not.toHaveProperty("password"); + expect(payload).not.toHaveProperty("url"); + expect(payload.host).toBe("localhost"); + }); + + it("should send a real new secret value the admin typed", () => { + const payload = buildCachePayload("node", { password: "brandnewpw" }, { forTesting: false }); + expect(payload.password).toBe("brandnewpw"); + }); +}); + +describe("secret handling", () => { + it("buildInitialValues never prefills a credential, even when the server reports it configured", () => { + const serverValues = { + host: "localhost", + password: REDACTED_VALUE, + url: REDACTED_VALUE, + sentinel_password: REDACTED_VALUE, + }; + const values = buildInitialValues(serverValues); + expect(values.password).toBe(""); + expect(values.url).toBe(""); + expect(values.sentinel_password).toBe(""); + // non-secret fields are still prefilled + expect(values.host).toBe("localhost"); + }); + + it("configuredSecretFields reports which credentials the server marked as set", () => { + const configured = configuredSecretFields({ + password: REDACTED_VALUE, + url: "", + host: "localhost", + }); + expect(configured.has("password")).toBe(true); + expect(configured.has("url")).toBe(false); + // a non-secret field is never reported as a configured secret + expect(configured.has("host")).toBe(false); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts index 088da21961c..7b9454a37c3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts @@ -1,4 +1,4 @@ -import { CACHE_FIELDS, CacheField, CacheSection, RedisType } from "./cacheSettingsFields"; +import { CACHE_FIELDS, CacheField, CacheSection, REDACTED_VALUE, RedisType } from "./cacheSettingsFields"; export type CacheFormValue = string | number | boolean | undefined; export type CacheFormValues = Record; @@ -11,7 +11,20 @@ export const isFieldVisible = (field: CacheField, redisType: RedisType): boolean export const fieldsForSection = (section: CacheSection, redisType: RedisType): CacheField[] => CACHE_FIELDS.filter((field) => field.section === section && isFieldVisible(field, redisType)); +const hasValue = (raw: unknown): boolean => raw !== undefined && raw !== null && raw !== ""; + +// Credential fields the server reports as configured (returned as the redacted +// marker). Used to show an "already set" hint without ever holding the secret. +export const configuredSecretFields = (currentValues: Record): ReadonlySet => + new Set(CACHE_FIELDS.filter((field) => field.secret && hasValue(currentValues[field.name])).map((f) => f.name)); + const initialValueForField = (field: CacheField, raw: unknown): CacheFormValue => { + // Never prefill a credential: the server sends the redacted marker for a + // configured secret, and echoing it back would persist the marker. + if (field.secret) { + return ""; + } + const source = raw ?? field.defaultValue; if (field.type === "boolean") { @@ -35,6 +48,11 @@ export const buildInitialValues = (currentValues: Record): Cach Object.fromEntries(CACHE_FIELDS.map((field) => [field.name, initialValueForField(field, currentValues[field.name])])); const saveValueForField = (field: CacheField, raw: CacheFormValue): CacheSavePayloadValue | undefined => { + // A redacted secret echoed back untouched must never be persisted as a value. + if (field.secret && raw === REDACTED_VALUE) { + return undefined; + } + if (field.type === "boolean") { return Boolean(raw); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.tsx index 4382769ae9c..fea2c04015b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.tsx @@ -8,7 +8,7 @@ import RedisTypeSelector from "./RedisTypeSelector"; import CacheFieldSection from "./CacheFieldSection"; import { EmbeddingModelOption } from "./CacheFormField"; import { REDIS_TYPES, REDIS_TYPE_DESCRIPTIONS, RedisType } from "./cacheSettingsFields"; -import { buildCachePayload, buildInitialValues, CacheFormValues } from "./cacheSettingsUtils"; +import { buildCachePayload, buildInitialValues, CacheFormValues, configuredSecretFields } from "./cacheSettingsUtils"; interface CacheSettingsProps { accessToken: string | null; @@ -25,6 +25,7 @@ const CacheSettings: React.FC = ({ accessToken }) => { const [embeddingModels, setEmbeddingModels] = useState([]); const [isTesting, setIsTesting] = useState(false); const [isSaving, setIsSaving] = useState(false); + const [configuredSecrets, setConfiguredSecrets] = useState>(new Set()); const loadCacheSettings = useCallback(async () => { if (!accessToken) { @@ -34,6 +35,7 @@ const CacheSettings: React.FC = ({ accessToken }) => { const data = (await getCacheSettingsCall(accessToken)) as { current_values?: Record }; const currentValues = data.current_values ?? {}; form.setFieldsValue(buildInitialValues(currentValues)); + setConfiguredSecrets(configuredSecretFields(currentValues)); setRedisType(toRedisType(currentValues.redis_type)); } catch (error) { console.error("Failed to load cache settings:", error); @@ -144,6 +146,7 @@ const CacheSettings: React.FC = ({ accessToken }) => { section="connection" redisType={redisType} embeddingModels={embeddingModels} + configuredSecrets={configuredSecrets} /> @@ -166,6 +169,7 @@ const CacheSettings: React.FC = ({ accessToken }) => { section="sentinel" redisType={redisType} embeddingModels={embeddingModels} + configuredSecrets={configuredSecrets} /> )} From 1aba849af2a7e53f83c6f876643c3378cb401933 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 21 Jul 2026 18:36:31 -0700 Subject: [PATCH 026/130] fix(ui): surface env-var-sourced theme and logging-callback settings (#34156) The UI theme and logging-callback read endpoints reported only stored config while the features resolve their values from the process environment, so a gateway configured purely through env vars showed blank settings pages even though branding rendered and callbacks fired. /get/ui_theme_settings read only litellm_settings.ui_theme_config; logo_url and favicon_url now fall back to UI_LOGO_PATH and LITELLM_FAVICON_URL when the stored config leaves them blank. process_callback (the logging-callbacks block of /get/config/callbacks) reported every callback env var as unset unless it lived in the config environment_variables overlay; it now falls back to os.getenv, matching the slack block. Secret values stay redacted for non-admins via the existing callback role gate. Stored values keep winning over the environment, so the UI-driven flow is unchanged. Resolves LIT-4667 --- litellm/proxy/common_utils/callback_utils.py | 8 +- .../proxy_setting_endpoints.py | 49 +++++++++- .../proxy/common_utils/test_callback_utils.py | 46 +++++++++ tests/test_litellm/proxy/test_proxy_server.py | 96 +++++++++++++++++++ .../test_proxy_setting_endpoints.py | 82 ++++++++++++++++ 5 files changed, 275 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c644ecc3dae..a9c2a12aff7 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,4 +1,5 @@ import copy +import os from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional import litellm @@ -564,11 +565,8 @@ def process_callback(_callback: str, callback_type: str, environment_variables: env_vars_dict: dict[str, str | None] = {} for _var in env_vars: - env_variable = environment_variables.get(_var, None) - if env_variable is None: - env_vars_dict[_var] = None - else: - env_vars_dict[_var] = env_variable + stored_value = environment_variables.get(_var, None) + env_vars_dict[_var] = stored_value if stored_value is not None else os.getenv(_var) return {"name": _callback, "variables": env_vars_dict, "type": callback_type} diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 48fa4bebaa3..42111cf17f2 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1,6 +1,8 @@ #### CRUD ENDPOINTS for UI Settings ##### import asyncio import json +import os +from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union from urllib.parse import urlparse @@ -35,6 +37,44 @@ _SSO_SENSITIVE_FIELDS: Set[str] = { "generic_client_secret", } +# Maps each UIThemeConfig field to the env var the UI branding path reads it +# from. /update/ui_theme_settings writes both the stored ui_theme_config and +# these env vars, so /get/ui_theme_settings resolves the same env vars to +# reflect a deployment branded purely through process env. +_UI_THEME_FIELD_ENV_VARS: dict[str, str] = { + "logo_url": "UI_LOGO_PATH", + "favicon_url": "LITELLM_FAVICON_URL", +} + + +def _is_public_http_url(value: str | None) -> bool: + """Whether a value is a plain http(s) URL with a host, safe to disclose publicly.""" + if not isinstance(value, str) or not value.strip(): + return False + parsed = urlparse(value.strip()) + return parsed.scheme in ("http", "https") and bool(parsed.netloc) + + +def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: + """Resolve one UI theme field to the value the branding path actually uses. + + The stored ui_theme_config wins; a field absent or blank there falls back to + the process environment. The branding path reads the env var, and stored + settings reach it by being pushed into the environment on save, so a value + supplied only as a process env var is live even though no stored entry exists. + + This endpoint is unauthenticated, so the env fallback only surfaces a public + http(s) URL: an operator can point UI_LOGO_PATH at a local filesystem path + (the branding path serves it server-side), and that path must not be + disclosed to anonymous callers. A stored value is already validated as a + public URL on write, so it passes through. + """ + stored = stored_values.get(field_name) + if isinstance(stored, str) and stored.strip(): + return stored + env_value = os.environ.get(_UI_THEME_FIELD_ENV_VARS[field_name]) + return env_value if _is_public_http_url(env_value) else None + class IPAddress(BaseModel): ip: str @@ -977,12 +1017,19 @@ async def get_ui_theme_settings(): # Load existing config config = await proxy_config.get_config() - return await _get_settings_with_schema( + result = await _get_settings_with_schema( settings_key="ui_theme_config", settings_class=UIThemeConfig, config=config, ) + stored_values = result.get("values", {}) + result["values"] = { + **stored_values, + **{field: _resolve_ui_theme_field(stored_values, field) for field in _UI_THEME_FIELD_ENV_VARS}, + } + return result + def _validate_public_image_url(value: Optional[str], field_name: str) -> None: """ diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 36ff3f3c399..8f390c096d7 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -84,6 +84,52 @@ def test_process_callback_with_no_required_env_vars(mock_get_env_vars): assert result["variables"] == {} +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"], +) +def test_process_callback_falls_back_to_process_env(mock_get_env_vars, monkeypatch): + """A callback env var set only in the process env must be surfaced. + + The logging integrations read their config from the process environment, so a + callback configured purely via env vars (IaC) is live even with no stored + entry. Reporting it as unset makes a working callback read as unconfigured. + """ + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "env-public-key") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "env-secret-key") + # stored config only carries the public key; the secret is env-only + environment_variables = {"LANGFUSE_PUBLIC_KEY": "db-public-key"} + + result = process_callback( + _callback="langfuse", + callback_type="success", + environment_variables=environment_variables, + ) + + # stored value wins; the env-only var is resolved rather than reported None + assert result["variables"] == { + "LANGFUSE_PUBLIC_KEY": "db-public-key", + "LANGFUSE_SECRET_KEY": "env-secret-key", + } + + +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_SECRET_KEY"], +) +def test_process_callback_reports_none_when_absent_everywhere(mock_get_env_vars, monkeypatch): + """A var set in neither the stored config nor the process env stays None.""" + monkeypatch.delenv("LANGFUSE_SECRET_KEY", raising=False) + + result = process_callback( + _callback="langfuse", + callback_type="success", + environment_variables={}, + ) + + assert result["variables"] == {"LANGFUSE_SECRET_KEY": None} + + def test_normalize_callback_names_none_returns_empty_list(): assert normalize_callback_names(None) == [] assert normalize_callback_names([]) == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 99a6c946a1e..8a19c6b4406 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1015,6 +1015,102 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): } +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"], +) +def test_get_config_callbacks_fall_back_to_process_env(mock_env_vars, monkeypatch): + """A callback configured purely via process env vars is surfaced. + + An IaC deployment sets LANGFUSE_* on the gateway and never touches the UI, + so nothing is stored in the config environment_variables overlay. The read + endpoint must still report the live values instead of blanks. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-env-only") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-env-only") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + + config_data = { + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": {}, + } + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + langfuse_cb = next( + (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None + ) + assert langfuse_cb is not None + assert langfuse_cb["variables"] == { + "LANGFUSE_PUBLIC_KEY": "pk-env-only", + "LANGFUSE_SECRET_KEY": "sk-env-only", + "LANGFUSE_HOST": "https://cloud.langfuse.com", + } + + +@patch( + "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", + return_value=["LANGFUSE_SECRET_KEY", "LANGFUSE_HOST"], +) +def test_get_config_callback_env_secrets_redacted_for_non_admin(mock_env_vars, monkeypatch): + """Surfacing env vars must not widen who can read secret values. + + The callback role gate redacts sensitive keys for anyone below full admin, + and that must hold whether the value came from the stored config or the + process env. A non-secret var (LANGFUSE_HOST) still resolves for context. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-env-only-secret") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + + config_data = { + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": {}, + } + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + langfuse_cb = next( + (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None + ) + assert langfuse_cb is not None + assert langfuse_cb["variables"]["LANGFUSE_SECRET_KEY"] == "REDACTED" + assert langfuse_cb["variables"]["LANGFUSE_HOST"] == "https://cloud.langfuse.com" + + def test_get_config_returns_email_settings(monkeypatch): """ Regression for https://github.com/BerriAI/litellm/issues/19221 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 671a6cd39ba..805baed9e1e 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 @@ -944,6 +944,88 @@ class TestProxySettingEndpoints: assert data["values"]["logo_url"] == "https://example.com/logo.png" assert data["values"]["favicon_url"] == "https://example.com/favicon.ico" + def test_get_ui_theme_settings_falls_back_to_process_env( + self, mock_proxy_config, monkeypatch + ): + """Branding supplied only as process env vars must surface in the read. + + A deployment that sets UI_LOGO_PATH / LITELLM_FAVICON_URL via IaC and + never touches the UI has no stored ui_theme_config, yet the branding is + live, so the settings page must reflect it rather than reading blank. + """ + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + monkeypatch.delenv("LITELLM_FAVICON_URL", raising=False) + monkeypatch.setenv("UI_LOGO_PATH", "https://cdn.example.com/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://cdn.example.com/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] == "https://cdn.example.com/logo.png" + assert values["favicon_url"] == "https://cdn.example.com/favicon.ico" + + def test_get_ui_theme_settings_stored_value_wins_over_env( + self, mock_auth, monkeypatch + ): + """A stored ui_theme_config field outranks the env var for that field. + + The env fallback only fills fields the stored config leaves blank, so the + UI-driven flow is unchanged while an unstored field still resolves. + """ + from litellm.proxy.proxy_server import proxy_config + + stored_config = { + "litellm_settings": { + "ui_theme_config": {"logo_url": "https://db.example.com/logo.png"} + } + } + + async def mock_get_config(): + return stored_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setenv("UI_LOGO_PATH", "https://env.example.com/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://env.example.com/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] == "https://db.example.com/logo.png" + assert values["favicon_url"] == "https://env.example.com/favicon.ico" + + def test_get_ui_theme_settings_reports_unset_when_absent_everywhere( + self, mock_proxy_config, monkeypatch + ): + """A field set in neither the stored config nor the env stays null.""" + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + monkeypatch.delenv("LITELLM_FAVICON_URL", raising=False) + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["logo_url"] is None + assert values["favicon_url"] is None + + def test_get_ui_theme_settings_does_not_disclose_local_path_env_value( + self, mock_proxy_config, monkeypatch + ): + """This endpoint is public, so an env-configured local filesystem branding + path must never be surfaced to anonymous callers; only public http(s) URLs. + """ + monkeypatch.setenv("UI_LOGO_PATH", "/mnt/secret/internal/logo.png") + monkeypatch.setenv("LITELLM_FAVICON_URL", "file:///etc/favicon.ico") + + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + values = response.json()["values"] + # the local path / file scheme is withheld rather than disclosed + assert values["logo_url"] is None + assert values["favicon_url"] is None + def test_get_ui_settings(self, mock_auth, monkeypatch): """Test retrieving UI settings with allowlist sanitization""" from unittest.mock import AsyncMock, MagicMock From a780d4e4e3208e387edea8bea276a961ed5c9c7d Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 21 Jul 2026 18:57:11 -0700 Subject: [PATCH 027/130] test(musty_leopard): cover customer chat/messages cost + streaming paths (#34164) * test(e2e): cover customer chat/messages cost + streaming paths Fills five uncovered P0 registry cells matching the customer's confirmed stack (OpenAI SDK, Bedrock, /v1/messages) and their per-request cost dependency: - /v1/messages logs cost that matches the x-litellm-response-cost header (LIT-4076) - OpenAI /chat/completions streams real content, and a non-streamed call is costed - Bedrock Converse /chat/completions returns real content non-streamed and streamed The streaming checks aggregate delta content and parse every chunk as JSON, so a clean-but-empty stream or a truncated chunk fails instead of passing on a bare 200. * test(e2e): add tool-use coverage for openai, bedrock converse, anthropic responses Function-calling regression guards on the paths the customer's agentic SDK usage exercises: OpenAI and Bedrock Converse /chat/completions, and Anthropic /v1/responses. The model is forced to call a weather tool and the test asserts the returned tool call names the function and carries JSON-parseable arguments with the expected field, so a dropped tool_call or malformed argument JSON fails instead of passing on a bare 200. Adds a minimal tool_calls field to the response OutMessage. * test(e2e): cover bedrock converse responses + thinking Adds llm.responses.bedrock_converse.basic/tool_use and llm.chat_completions.bedrock_converse.thinking. The thinking test enables extended thinking and requires reasoning_content plus a real answer, so a path that drops the reasoning block fails rather than passing. * test(e2e): cover bedrock embeddings + openai structured output and reasoning Bedrock Titan embeddings return a real vector; OpenAI structured output must yield schema-conforming JSON with the correct extracted values (age==42, not just valid JSON); an OpenAI reasoning call must report reasoning tokens, so a non-reasoning fallback fails. Adds response_format to ChatBody and reasoning-token details to Usage. * test(e2e): cover vision + streaming tool calls on openai and bedrock converse Vision on both providers must describe the image (not just 200); the streamed OpenAI tool call is reassembled from its fragments and its argument JSON parsed, so a stream that never completes the call or splits its JSON fails. Extends ChatMessage content to a typed text/image union. * test(e2e): cover openai prompt caching hit on repeated large prefix A repeated large-prefix prompt must report cached prompt tokens on the second call, so a cache regression that stops reusing the prefix (and silently re-bills full input) fails here. * test(e2e): cover openai audio speech + bedrock rerank and image generation Marks the OpenAI TTS cell and adds Bedrock Titan rerank (top_n honored, scored) and Bedrock Titan image generation (returns b64/url), the customer's non-chat AWS surfaces. * test(e2e): cover end-user (customer) create persistence mgmt.end_user.new.happy_path: create an end-user via /customer/new and confirm /customer/info reports it, the end-user-identity surface the customer relies on for per-customer controls. Adds customer models + management-client methods. * test(e2e): enforce key model allow-list on the passthrough route other.auth.passthrough.model_allowlist_enforced: a key scoped to gemini must be denied a claude call through the anthropic passthrough route (403), so custom-auth scoping is not bypassable by going through passthrough instead of /chat/completions. * test(e2e): address Greptile - assert stream data events, correlate messages spend by key - streaming: assert len(stream_events) > 1 instead of chunks > 1, since chunks counts the terminal data: [DONE] marker and would pass a single content event - messages cost: correlate the spend row by the unique scoped key rather than the Anthropic response id, which need not equal the proxy spend-log request_id --- .../test_chat_completions_regression_e2e.py | 551 +++++++++++++++++- .../test_image_generation_e2e.py | 38 +- .../e2e/llm_translation/test_messages_e2e.py | 64 +- .../llm_translation/test_passthrough_e2e.py | 25 +- tests/e2e/llm_translation/test_rerank_e2e.py | 43 +- .../e2e/llm_translation/test_responses_e2e.py | 100 +++- tests/e2e/management/management_client.py | 33 ++ tests/e2e/management/test_management_e2e.py | 15 + tests/e2e/models.py | 47 +- 9 files changed, 894 insertions(+), 22 deletions(-) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 13992744f42..8d3622e441a 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -19,17 +19,176 @@ from __future__ import annotations import os import pytest +from pydantic import BaseModel from e2e_config import require_env, unique_marker -from e2e_http import unwrap +from e2e_http import StreamingResponse, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, LiteLLMParamsBody +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + ImageContentPart, + ImageUrl, + LiteLLMParamsBody, + TextContentPart, + ThinkingParam, +) from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e COHERE_BACKEND = "cohere/command-r-08-2024" GEMINI_BACKEND = "gemini/gemini-2.5-flash" +OPENAI_BACKEND = "openai/gpt-5.6" +BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +class _StreamToolCallFunction(BaseModel): + name: str | None = None + arguments: str | None = None + + +class _StreamToolCall(BaseModel): + function: _StreamToolCallFunction = _StreamToolCallFunction() + + +class _StreamDelta(BaseModel): + content: str | None = None + tool_calls: list[_StreamToolCall] | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta = _StreamDelta() + + +class _StreamChunk(BaseModel): + choices: list[_StreamChoice] = [] + + +def _streamed_tool_call(events: list[str]) -> tuple[str, str]: + """Reassemble the tool call streamed across chunks: the name arrives once and the + arguments arrive as fragments, so concatenating both and parsing the arguments as + JSON catches a stream that never completes the call or splits its argument JSON.""" + chunks = [_StreamChunk.model_validate_json(event) for event in events] + calls = [call for chunk in chunks for choice in chunk.choices for call in (choice.delta.tool_calls or [])] + name = "".join(call.function.name or "" for call in calls) + arguments = "".join(call.function.arguments or "" for call in calls) + return name, arguments + + +CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" +OPENAI_VISION_BACKEND = "openai/gpt-4o" + +# OpenAI caches a shared prompt prefix once it exceeds ~1024 tokens; this is well +# past that, so a repeat call reports cached prompt tokens. +CACHE_PREFIX = ( + "You are a meticulous assistant. Follow these standing instructions exactly. " + * 300 +) + + +def _vision_messages() -> list[ChatMessage]: + return [ + ChatMessage( + role="user", + content=[ + TextContentPart(text="What animal is in this image? Answer in one word."), + ImageContentPart(image_url=ImageUrl(url=CAT_IMAGE_URL)), + ], + ) + ] + + +def _assert_describes_cat(response: ChatResponse) -> None: + assert response.choices, f"vision returned no choices: {response}" + message = response.choices[0].message + content = (message.content if message else None) or "" + assert "cat" in content.lower() or "feline" in content.lower(), ( + f"vision response did not describe the image: {content[:200]}" + ) + + +def _streamed_text(events: list[str]) -> str: + """Concatenate the delta content across streamed chunks. Parsing every event as + JSON also fails loudly on a truncated or garbled chunk (the vertex/gemini image + streaming regression class), so an incomplete stream cannot pass as content.""" + chunks = [_StreamChunk.model_validate_json(event) for event in events] + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +def _assert_streamed_completion(result: StreamingResponse) -> None: + """A streamed /chat/completions must deliver real content, not a clean-but-empty + stream (the #28991 class on the streaming path).""" + assert result.ok and result.is_streaming, f"stream was not established: {result}" + assert result.stream_error is None, f"stream carried an error event: {result.stream_error}" + assert len(result.stream_events) > 1, f"stream did not deliver multiple data events: {result}" + assert _streamed_text(result.stream_events).strip(), ( + f"stream completed with no content deltas: {result.stream_events[:3]}" + ) + + +def _bedrock_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ) + + +class _WeatherArgs(BaseModel): + location: str + + +_WEATHER_TOOL = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a location", + parameters={ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + ) +) + + +def _assert_weather_tool_call(response: ChatResponse) -> None: + """The model, forced to call the tool, must return a get_weather call whose + arguments parse as JSON and carry a location. A regression that drops tool_calls + or emits malformed argument JSON fails here rather than passing on a 200.""" + assert response.choices, f"chat returned no choices: {response}" + message = response.choices[0].message + calls = message.tool_calls if message else None + assert calls, f"model returned no tool call for a tool-forced prompt: {response}" + weather = next((call for call in calls if call.function.name == "get_weather"), None) + assert weather is not None, f"expected a get_weather call, got {[c.function.name for c in calls]}" + assert weather.function.arguments, f"get_weather call carried no arguments: {weather}" + args = _WeatherArgs.model_validate_json(weather.function.arguments) + assert args.location.strip(), f"get_weather arguments missing location: {weather.function.arguments}" + + +class _Person(BaseModel): + name: str + age: int + + +_PERSON_SCHEMA: dict[str, object] = { + "type": "json_schema", + "json_schema": { + "name": "person", + "strict": True, + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}, + "required": ["name", "age"], + "additionalProperties": False, + }, + }, +} CHAT_MODELS: tuple[tuple[str, str], ...] = ( ("gpt-5.5", "openai"), @@ -219,3 +378,391 @@ class TestHostedVllmChat: assert response.choices, f"hosted_vllm chat returned no choices: {response}" content = response.choices[0].message.content if response.choices[0].message else None assert content and content.strip(), f"hosted_vllm empty content: {response}" + + +class TestOpenAIChatCompletions: + """OpenAI /chat/completions, the SDK path the customer runs against the proxy. + + The streamed call must deliver real content deltas (a clean-but-empty stream is + the regression), and a non-streamed call must be costed so per-request spend and + the response-cost header stay accurate. + """ + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.openai.basic.nonstream.cost_logged", + exercised_on=["chat_completions"], + ) + def test_openai_chat_logs_cost( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-cost-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=16, + ), + ) + ) + assert response.choices, f"openai chat returned no choices: {response}" + + rows = client.proxy.poll_logs_for_key( + key, min_rows=1, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + priced = [r for r in rows if (r.spend or 0) > 0] + assert priced, f"openai chat was not costed on key ...{key[-6:]}: {rows}" + assert priced[0].status == "success", f"openai chat spend status={priced[0].status!r}" + + @pytest.mark.covers( + "llm.chat_completions.openai.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-tool-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.openai.structured_output.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_structured_output_conforms_to_schema( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-schema-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="Extract the person. John Doe is 42 years old.")], + response_format=_PERSON_SCHEMA, + max_tokens=128, + ), + ) + ) + assert response.choices, f"structured output returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content, f"structured output returned empty content: {response}" + person = _Person.model_validate_json(content) + assert person.name.strip() and person.age == 42, ( + f"schema-constrained extraction was wrong: {person}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.thinking.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_reasoning_reports_reasoning_tokens( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-reasoning-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content="A train travels 60 miles in 1.5 hours. What is its average speed in mph?", + ) + ], + reasoning_effort="low", + max_tokens=2048, + ), + ) + ) + assert response.choices, f"reasoning call returned no choices: {response}" + message = response.choices[0].message + assert message and message.content and message.content.strip(), f"reasoning call had no answer: {response}" + details = response.usage.completion_tokens_details if response.usage else None + assert details and details.reasoning_tokens and details.reasoning_tokens > 0, ( + f"a reasoning model must report reasoning tokens, got usage={response.usage}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-vision-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_VISION_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) + + @pytest.mark.covers( + "llm.chat_completions.openai.prompt_cache_5m.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_prompt_cache_hits_on_repeat( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + body = ChatBody( + model=model, + messages=[ + ChatMessage(role="system", content=CACHE_PREFIX), + ChatMessage(role="user", content="Reply with the single word pong."), + ], + max_tokens=16, + ) + unwrap(client.proxy.chat(key, body)) + second = unwrap(client.proxy.chat(key, body)) + + details = second.usage.prompt_tokens_details if second.usage else None + assert details and details.cached_tokens and details.cached_tokens > 0, ( + f"a repeated large-prefix prompt must report cached prompt tokens, got usage={second.usage}" + ) + + @pytest.mark.covers( + "llm.chat_completions.openai.tool_use.stream.works", + exercised_on=["chat_completions"], + ) + def test_openai_chat_streams_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + require_env("OPENAI_API_KEY") + model = f"e2e-openai-tool-stream-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + stream=True, + ), + ) + assert result.ok and result.is_streaming, f"tool stream was not established: {result}" + assert result.stream_error is None, f"tool stream carried an error event: {result.stream_error}" + name, arguments = _streamed_tool_call(result.stream_events) + assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}" + args = _WeatherArgs.model_validate_json(arguments) + assert args.location.strip(), f"streamed tool call arguments missing location: {arguments!r}" + + +class TestBedrockConverseChatCompletions: + """Bedrock Converse via /chat/completions, the customer's AWS stack. A non-OpenAI + provider must return real content on both the non-streamed and streamed paths. + """ + + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model(model, _bedrock_params()) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-chat") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=32, + ), + ) + ) + assert response.choices, f"bedrock converse chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"bedrock converse returned empty content: {response}" + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.thinking.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_returns_reasoning( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-thinking") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="What is 17 times 23? Think it through step by step.")], + thinking=ThinkingParam(type="enabled", budget_tokens=1024), + max_tokens=2048, + ), + ) + ) + assert response.choices, f"bedrock thinking returned no choices: {response}" + message = response.choices[0].message + assert message and message.content and message.content.strip(), ( + f"bedrock thinking returned no answer content: {response}" + ) + assert message.reasoning_content and message.reasoning_content.strip(), ( + "thinking was enabled but no reasoning_content came back on the Bedrock Converse path" + ) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_bedrock_converse_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-bedrock-vision") + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index d4080afb7dc..1ba78a7e083 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -9,7 +9,7 @@ from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, ImagesResult from lifecycle import ResourceManager @@ -18,6 +18,15 @@ from models import LiteLLMParamsBody pytestmark = pytest.mark.e2e +def _assert_image_returned(body: str) -> None: + parsed = ImagesResult.model_validate_json(body) + assert parsed.data, f"/images/generations returned no data: {body[:300]}" + first = parsed.data[0] + assert first.b64_json or first.url, ( + f"generated image has neither b64_json nor url: {body[:300]}" + ) + + class TestImageGeneration: @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( @@ -35,9 +44,26 @@ class TestImageGeneration: result = endpoints_client.images(key, model, "Draw a cute cat") require_successful_call(result) - parsed = ImagesResult.model_validate_json(result.body) - assert parsed.data, f"/images/generations returned no data: {result.body[:300]}" - first = parsed.data[0] - assert first.b64_json or first.url, ( - f"generated image has neither b64_json nor url: {result.body[:300]}" + _assert_image_returned(result.body) + + @pytest.mark.covers("llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"]) + def test_bedrock_image_generation_returns_image( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + model = f"e2e-bedrock-image-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.titan-image-generator-v2:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.images(key, model, "Draw a cute cat") + require_successful_call(result) + _assert_image_returned(result.body) diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index a14bf8e82b3..44376218c6b 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -10,7 +10,7 @@ from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from e2e_http import require_successful_call, unwrap from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager @@ -20,11 +20,14 @@ from models import ( ChatMessage, JsonSchemaProperty, LiteLLMParamsBody, + SpendLogRow, ToolInputSchema, ) pytestmark = pytest.mark.e2e +ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" + WEATHER_TOOL = AnthropicCustomTool( name="get_weather", description="Get the current weather for a city.", @@ -35,6 +38,11 @@ WEATHER_TOOL = AnthropicCustomTool( ) +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + class TestAnthropicMessages: def _register( self, endpoints_client: EndpointsClient, resources: ResourceManager @@ -43,7 +51,7 @@ class TestAnthropicMessages: model_id = endpoints_client.create_model( model, LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" ), ) resources.defer(lambda: endpoints_client.delete_model(model_id)) @@ -61,6 +69,58 @@ class TestAnthropicMessages: assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged") + def test_messages_logs_cost_matching_the_response_header( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + require_env("ANTHROPIC_API_KEY") + model = f"e2e-messages-cost-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}") + require_successful_call(result) + parsed = MessagesResult.model_validate_json(result.body) + assert parsed.role == "assistant" and parsed.text.strip(), ( + f"/v1/messages returned no assistant text: {result.body[:300]}" + ) + + # The customer reads per-request cost off the response header (LIT-4076), so + # it must be present and positive on /v1/messages, not only /chat/completions. + header_cost = result.response_cost + assert header_cost is not None and header_cost > 0, ( + "x-litellm-response-cost header missing or non-positive on /v1/messages; " + f"headers={result.headers}" + ) + + # Correlate the spend row by the unique scoped key, not the Anthropic response + # id: on /v1/messages the spend-log request_id is the proxy's own call id, which + # need not equal the message body id, so an id-based poll can miss a correctly + # logged row and time out. The key is fresh per test, so its only priced row is + # this call. + def _priced(rows: list[SpendLogRow]) -> bool: + return any(r.spend is not None and r.spend > 0 for r in rows) + + rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced) + priced = [r for r in rows if r.spend is not None and r.spend > 0] + assert priced, ( + f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" + ) + row = priced[0] + assert (row.prompt_tokens or 0) > 0 and (row.completion_tokens or 0) > 0, ( + f"messages spend row missing token counts, so the cost is not real usage: {row}" + ) + assert row.spend is not None and _approx_equal(row.spend, header_cost), ( + f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}; " + "the customer bills against the header, so the two must match" + ) + @pytest.mark.covers("llm.messages.anthropic.basic.stream.works") def test_messages_streams_completion( self, endpoints_client: EndpointsClient, resources: ResourceManager diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index c8806faf3ea..ed5c657d23e 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -15,7 +15,8 @@ import pytest from e2e_config import unique_marker from e2e_http import StreamingResponse, require_successful_call -from models import SpendLogRow +from lifecycle import ResourceManager +from models import KeyGenerateBody, SpendLogRow from passthrough_client import ( AnthropicTool, GeminiFunctionDeclaration, @@ -157,3 +158,25 @@ def test_anthropic_passthrough_tool_call_logs_cost( row = _fetch_cost_breakdown(client, result) assert row.custom_llm_provider == "anthropic" + + +class TestPassthroughModelAllowlist: + """A passthrough route must honor the calling key's model allow-list. + + The customer fronts native provider calls through the proxy with custom auth, + so a key scoped to one model must not reach a different model just because the + request goes through the passthrough route rather than /chat/completions. + """ + + @pytest.mark.covers("other.auth.passthrough.model_allowlist_enforced") + def test_passthrough_denies_model_outside_key_allowlist( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + key = client.proxy.generate_key(KeyGenerateBody(models=["gemini-2.5-flash"])) + resources.defer(lambda: client.proxy.delete_key(key)) + + result = client.anthropic_message(key, "claude-haiku-4-5", f"say hi {unique_marker()}") + assert result.status_code == 403, ( + "a key restricted to gemini-2.5-flash must be denied a claude passthrough call, " + f"got {result.status_code}: {result.body[:300]}" + ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 31801b306a8..0857ff65a52 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -9,7 +9,7 @@ from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from e2e_http import require_successful_call from endpoints_client import EndpointsClient, RerankResult from lifecycle import ResourceManager @@ -23,6 +23,16 @@ DOCUMENTS = [ "Washington, D.C. is the capital of the United States.", "Capital punishment has existed in the United States since before it was a country.", ] +QUERY = "What is the capital of the United States?" + + +def _assert_top_n_scored(body: str) -> None: + parsed = RerankResult.model_validate_json(body) + assert parsed.results, f"/rerank returned no results: {body[:300]}" + assert len(parsed.results) <= 3, f"top_n=3 not honored: {body[:300]}" + assert parsed.results[0].relevance_score is not None, ( + f"top rerank result has no relevance_score: {body[:300]}" + ) class TestRerank: @@ -38,13 +48,28 @@ class TestRerank: resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank( - key, model, "What is the capital of the United States?", DOCUMENTS, top_n=3 - ) + result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) require_successful_call(result) - parsed = RerankResult.model_validate_json(result.body) - assert parsed.results, f"/rerank returned no results: {result.body[:300]}" - assert len(parsed.results) <= 3, f"top_n=3 not honored: {result.body[:300]}" - assert parsed.results[0].relevance_score is not None, ( - f"top rerank result has no relevance_score: {result.body[:300]}" + _assert_top_n_scored(result.body) + + @pytest.mark.covers("llm.rerank.bedrock.basic.nonstream.works", exercised_on=["rerank"]) + def test_bedrock_rerank_scores_top_n( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + model = f"e2e-bedrock-rerank-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="bedrock/amazon.rerank-v1:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) + require_successful_call(result) + _assert_top_n_scored(result.body) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index bd98f11c045..d24d2b53b71 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -13,7 +13,7 @@ from typing import cast import pytest from pydantic import BaseModel, ValidationError -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from e2e_http import require_successful_call from endpoints_client import ( EndpointsClient, @@ -29,6 +29,26 @@ from models import LiteLLMParamsBody pytestmark = pytest.mark.e2e +BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + +WEATHER_TOOL = ResponsesFunctionTool( + name="get_weather", + description="Get the weather for a location", + parameters=FunctionParameters( + properties={"location": FunctionParameterProperty(type="string")}, + required=["location"], + ), +) + + +def _bedrock_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=BEDROCK_CONVERSE_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ) + class WeatherArguments(BaseModel): location: str @@ -190,6 +210,84 @@ class TestResponses: parsed = ResponsesResult.model_validate_json(result.body) assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + @pytest.mark.covers("llm.responses.anthropic.tool_use.nonstream.works") + def test_responses_anthropic_returns_function_call( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses_with_tools( + key, + model, + "What is the weather in San Francisco? Use the get_weather tool.", + [ + ResponsesFunctionTool( + name="get_weather", + description="Get the weather for a location", + parameters=FunctionParameters( + properties={"location": FunctionParameterProperty(type="string")}, + required=["location"], + ), + ) + ], + ) + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + function_call = next( + (call for call in parsed.function_calls if call.name == "get_weather"), + None, + ) + assert function_call is not None, f"no get_weather function call: {result.body[:500]}" + assert function_call.arguments is not None + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + + @pytest.mark.covers("llm.responses.bedrock_converse.basic.nonstream.works") + def test_responses_bedrock_returns_completion( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model(model, _bedrock_params()) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses(key, model, "reply with one word") + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + assert parsed.text.strip(), f"/responses over bedrock returned no output text: {result.body[:300]}" + + @pytest.mark.covers("llm.responses.bedrock_converse.tool_use.nonstream.works") + def test_responses_bedrock_returns_function_call( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + model = f"e2e-responses-{unique_marker()}" + model_id = endpoints_client.create_model(model, _bedrock_params()) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.responses_with_tools( + key, model, "What is the weather in San Francisco? Use the get_weather tool.", [WEATHER_TOOL] + ) + require_successful_call(result) + parsed = ResponsesResult.model_validate_json(result.body) + function_call = next((call for call in parsed.function_calls if call.name == "get_weather"), None) + assert function_call is not None, f"no get_weather function call over bedrock: {result.body[:500]}" + assert function_call.arguments is not None + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + def _parse_stream_event( event: str, diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index a9dedac8e61..cdc31aeea79 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -14,6 +14,10 @@ from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, Un from models import ( ChatBody, ChatMessage, + CustomerDeleteBody, + CustomerInfoParams, + CustomerNewBody, + CustomerResponse, KeyBlockBody, KeyDeleteBody, KeyGenerateBody, @@ -270,6 +274,35 @@ class ManagementClient: ) ).user_id + def create_customer(self, user_id: str) -> str: + _ = unwrap( + self.proxy.transport.post( + "/customer/new", + headers=self.proxy.transport.master, + json=CustomerNewBody(user_id=user_id), + response_type=CustomerResponse, + ) + ) + return user_id + + def customer_info(self, end_user_id: str) -> CustomerResponse: + return unwrap( + self.proxy.transport.get( + "/customer/info", + headers=self.proxy.transport.master, + params=CustomerInfoParams(end_user_id=end_user_id), + response_type=CustomerResponse, + ) + ) + + def delete_customer(self, user_id: str) -> None: + _ = self.proxy.transport.post( + "/customer/delete", + headers=self.proxy.transport.master, + json=CustomerDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + def update_user(self, body: UserUpdateBody) -> None: _ = unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 18bc384a879..9b398963ac9 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -610,3 +610,18 @@ class TestManagementRoutePermissions: f"/team/info returned {team_probe.status_code}: {team_probe.body[:300]}" ) assert client.user_count(user_id) == 0, f"user {user_id} was created despite the 403 route denial" + + +class TestCustomer: + @pytest.mark.covers("mgmt.end_user.new.happy_path") + def test_customer_create_persists_to_info( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + customer = f"e2e-customer-{unique_marker()}" + client.create_customer(customer) + resources.defer(lambda: client.delete_customer(customer)) + + info = client.customer_info(customer) + assert info.user_id == customer, ( + f"/customer/info did not report the created end-user; got {info.user_id!r}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8b6c454fa1e..b3ea9346180 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -115,6 +115,18 @@ class KeyInfoResponse(BaseModel): # ---------- customers ---------- +class CustomerNewBody(BaseModel): + user_id: str + + +class CustomerResponse(BaseModel): + user_id: str | None = None + + +class CustomerInfoParams(BaseModel): + end_user_id: str + + class CustomerDeleteBody(BaseModel): user_ids: list[str] @@ -126,9 +138,26 @@ class ChatMetadata(BaseModel): tags: list[str] | None = None +class ImageUrl(BaseModel): + url: str + + +class TextContentPart(BaseModel): + type: str = "text" + text: str + + +class ImageContentPart(BaseModel): + type: str = "image_url" + image_url: ImageUrl + + +ContentPart = TextContentPart | ImageContentPart + + class ChatMessage(BaseModel): role: str - content: str + content: str | list[ContentPart] class CacheControl(BaseModel): @@ -180,6 +209,7 @@ class ChatBody(BaseModel): tools: list[ChatTool] | None = None tool_choice: str | None = None guardrails: list[str] | None = None + response_format: dict[str, object] | None = None class RouterSettingsOverride(BaseModel): @@ -203,9 +233,19 @@ class ReliabilityChatBody(ChatBody): router_settings_override: RouterSettingsOverride | None = None +class ToolCallFunction(BaseModel): + name: str | None = None + arguments: str | None = None + + +class ToolCall(BaseModel): + function: ToolCallFunction = ToolCallFunction() + + class OutMessage(BaseModel): content: str | None = None reasoning_content: str | None = None + tool_calls: list[ToolCall] | None = None class ChatChoice(BaseModel): @@ -216,6 +256,10 @@ class PromptTokensDetails(BaseModel): cached_tokens: int | None = None +class CompletionTokensDetails(BaseModel): + reasoning_tokens: int | None = None + + class Usage(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None @@ -223,6 +267,7 @@ class Usage(BaseModel): cache_read_input_tokens: int | None = None cache_creation_input_tokens: int | None = None prompt_tokens_details: PromptTokensDetails | None = None + completion_tokens_details: CompletionTokensDetails | None = None class ChatResponse(BaseModel): From 33fb38056ddc16224a6ec5e6939167f6f87194d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:17:11 -0700 Subject: [PATCH 028/130] fix(e2e): register the load mock model through load_key instead of an autouse fixture so harness tests run without a proxy --- tests/e2e/load/conftest.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/e2e/load/conftest.py b/tests/e2e/load/conftest.py index d9b1b2d9d69..89a571af83e 100644 --- a/tests/e2e/load/conftest.py +++ b/tests/e2e/load/conftest.py @@ -51,10 +51,8 @@ def _model_is_servable(proxy: ProxyClient, model_name: str) -> bool: return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data) -@pytest.fixture(scope="session", autouse=True) -def _ensure_load_model( # pyright: ignore[reportUnusedFunction] # pytest autouse session fixture, wired by name - client: LoadClient, -) -> Iterator[None]: +@pytest.fixture(scope="session") +def ensure_load_model(client: LoadClient) -> Iterator[None]: proxy = client.proxy if _model_is_servable(proxy, LOAD_MODEL): yield @@ -78,7 +76,9 @@ def _ensure_load_model( # pyright: ignore[reportUnusedFunction] # pytest autou @pytest.fixture -def load_key(resources: ResourceManager, client: LoadClient) -> str: +def load_key( + resources: ResourceManager, client: LoadClient, ensure_load_model: None +) -> str: key = client.proxy.generate_key(KeyGenerateBody(models=[LOAD_MODEL], user_id="e2e-load")) resources.defer(lambda: client.proxy.delete_key(key)) return key From f3f89d6177173ce060dd9477f44851b72982ff0e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 21 Jul 2026 20:11:49 -0700 Subject: [PATCH 029/130] fix(proxy): raise dashboard session budget default to $1 and make it configurable in config and Admin UI Every dashboard login mints a 24h session key whose max_budget comes from litellm.max_ui_session_budget, and all dashboard LLM traffic (playground, auto router per-tier Test Connection probes) spends against and is gated by that one key. The $0.25 default locked sessions out mid-testing with "Budget has been exceeded ... Max budget: 0.25" and the setting appeared in no docs, no UI, and no error text, so it read as a hardcoded cap. Raise the default to $1. Give the setting an explicit typed arm in the config loader (float coercion for env-var strings, null disables the cap). Surface it on the Admin UI General settings tab through the existing litellm_settings bridge as a new Dollar field type (positive USD, unbounded above; the existing Float type is validated to (0, 1] for fractions), with a spec-level default so clearing the field restores $1 instead of silently removing the cap, and enroll it in LITELLM_SETTINGS_SAFE_DB_OVERRIDES so UI edits propagate to peer workers. --- litellm/__init__.py | 4 +- litellm/constants.py | 1 + litellm/proxy/proxy_server.py | 36 ++++- tests/test_litellm/proxy/test_proxy_server.py | 142 ++++++++++++++++++ .../_components/general_settings.tsx | 11 ++ 5 files changed, 187 insertions(+), 7 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 2f6643c644c..55821012df9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -427,7 +427,9 @@ default_team_settings: Optional[List] = None max_user_budget: Optional[float] = None default_max_internal_user_budget: Optional[float] = None max_internal_user_budget: Optional[float] = None -max_ui_session_budget: Optional[float] = 0.25 # $0.25 USD budgets for UI Chat sessions +max_ui_session_budget: Optional[float] = ( + 1.0 # USD budget for each dashboard login session (playground, test connection) +) internal_user_budget_duration: Optional[str] = None tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None max_end_user_budget: Optional[float] = None diff --git a/litellm/constants.py b/litellm/constants.py index 05944c81ea2..94d3e0b2b66 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1524,6 +1524,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ # test_general_settings_ui_fields_are_db_overridable enforces that pairing. "enable_anthropic_prompt_caching", "anthropic_prompt_caching_ttl", + "max_ui_session_budget", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b40abed19e..dda311ef0b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4569,6 +4569,11 @@ class ProxyConfig: verbose_proxy_logger.debug( f"{blue_color_code} Initialized polling via cache: enabled={polling_via_cache_enabled}, native_background_mode={native_background_mode}, ttl={polling_cache_ttl}{reset_color_code}" ) + elif key == "max_ui_session_budget": + litellm.max_ui_session_budget = float(value) if value is not None else None + verbose_proxy_logger.debug( + f"{blue_color_code} setting litellm.max_ui_session_budget={litellm.max_ui_session_budget}{reset_color_code}" + ) elif key == "default_team_settings": for idx, team_setting in enumerate(value): # run through pydantic validation try: @@ -14845,10 +14850,11 @@ GeneralSettingsUILiteLLMValue = Union[float, bool, str, None] class GeneralSettingsUILiteLLMFieldSpec(TypedDict): - type: Literal["Float", "Boolean", "Select"] + type: Literal["Float", "Dollar", "Boolean", "Select"] description: str options: NotRequired[tuple[str, ...]] tab: NotRequired[str] # Admin UI sub-tab this field renders under; None groups it with the rest + default: NotRequired[float] # reset/clear restores this instead of None; fields whose None means fail-open set it _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec] = { @@ -14874,21 +14880,32 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec "tab": "prompt_caching", "description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.", }, + "max_ui_session_budget": { + "type": "Dollar", + "default": 1.0, + "description": ( + "USD spend cap for each dashboard login session; covers LLM calls made from the dashboard " + "such as the playground and auto router Test Connection. Each login starts a fresh session " + "with this budget. Clearing restores the $1 default." + ), + }, } def _general_settings_ui_litellm_default( - field_type: Literal["Float", "Boolean", "Select"], + spec: GeneralSettingsUILiteLLMFieldSpec, ) -> GeneralSettingsUILiteLLMValue: """The value a field falls back to when it is cleared or reset.""" - return False if field_type == "Boolean" else None + if "default" in spec: + return spec["default"] + return False if spec["type"] == "Boolean" else None def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue: spec = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name] field_type = spec["type"] if value is None or value == "": - return _general_settings_ui_litellm_default(field_type) + return _general_settings_ui_litellm_default(spec) match field_type: case "Boolean": if not isinstance(value, bool): @@ -14912,6 +14929,13 @@ def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, ) return float(value) + case "Dollar": + if isinstance(value, bool) or not isinstance(value, (int, float)) or float(value) <= 0: + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be a positive dollar amount or empty"}, + ) + return float(value) case _: assert_never(field_type) @@ -14934,7 +14958,7 @@ async def _persist_general_settings_ui_litellm_field( async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict: config = await proxy_config.get_config() before_value = config.get("litellm_settings", {}).get(field_name) - default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]["type"]) + default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]) setattr(litellm, field_name, default_value) if "litellm_settings" in config: config["litellm_settings"].pop(field_name, None) @@ -15108,7 +15132,7 @@ async def get_config_list( ) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) - default_value = _general_settings_ui_litellm_default(spec["type"]) + default_value = _general_settings_ui_litellm_default(spec) stored_in_db_litellm: Optional[bool] if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a100e7837f4..8624c98da52 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2588,6 +2588,69 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp litellm.max_budget = original_max_budget +def test_max_ui_session_budget_default_is_one_dollar(): + """LIT-4662: the dashboard session budget default is a product decision; the + old 0.25 default locked admins out of auto router Test Connection and the + playground mid-session with an error that looked like a hardcoded cap.""" + assert litellm.max_ui_session_budget == 1.0 + + +@pytest.mark.asyncio +async def test_load_config_max_ui_session_budget_applied_and_coerced(tmp_path, monkeypatch): + """ + max_ui_session_budget configured via os.environ resolves to a string; + load_config must coerce it to float so every dashboard session key is + minted with a numeric max_budget. + """ + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("UI_SESSION_BUDGET", "2.5") + test_config = { + "model_list": [], + "litellm_settings": {"max_ui_session_budget": "os.environ/UI_SESSION_BUDGET"}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_budget = litellm.max_ui_session_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert isinstance(litellm.max_ui_session_budget, float) + assert litellm.max_ui_session_budget == 2.5 + finally: + litellm.max_ui_session_budget = original_budget + + +@pytest.mark.asyncio +async def test_load_config_max_ui_session_budget_none_disables_cap(tmp_path): + """ + max_ui_session_budget: null in config disables the dashboard session cap + entirely (session keys minted with no max_budget); load_config must pass + None through instead of raising on float(None). + """ + from litellm.proxy.proxy_server import ProxyConfig + + test_config = { + "model_list": [], + "litellm_settings": {"max_ui_session_budget": None}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_budget = litellm.max_ui_session_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert litellm.max_ui_session_budget is None + finally: + litellm.max_ui_session_budget = original_budget + + @pytest.mark.asyncio async def test_load_config_default_internal_user_params_max_budget_scientific_notation(tmp_path): """ @@ -9048,6 +9111,85 @@ def test_general_settings_ui_fields_are_db_overridable(): ) +@pytest.mark.asyncio +async def test_update_config_field_max_ui_session_budget_sets_live_value(monkeypatch): + """LIT-4662: the dashboard session budget is editable from the Admin UI General tab. + A Dollar field must accept values above 1 (the old Float type capped at 1, which cannot + express a dollar budget), apply live via setattr, and persist under litellm_settings.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, "max_ui_session_budget", 1.0) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="max_ui_session_budget", + field_value=25.0, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert litellm.max_ui_session_budget == 25.0 + assert saved["litellm_settings"]["max_ui_session_budget"] == 25.0 + + +@pytest.mark.parametrize("bad_value", [True, "abc", -1, 0, [2.5]]) +def test_validate_max_ui_session_budget_rejects_malformed(bad_value): + """A Dollar field accepts only positive numbers; zero would block every dashboard + LLM call at mint and non-numerics would break session key generation.""" + from fastapi import HTTPException + + from litellm.proxy.proxy_server import _validate_general_settings_ui_litellm_value + + with pytest.raises(HTTPException) as exc_info: + _validate_general_settings_ui_litellm_value("max_ui_session_budget", bad_value) + assert exc_info.value.status_code == 400 + + +@pytest.mark.parametrize("empty_value", [None, ""]) +def test_validate_max_ui_session_budget_empty_restores_default(empty_value): + """Clearing the field in the UI restores the shipped $1 default rather than None; + None would silently remove the session spend guardrail (unlimited budget), which + must stay a deliberate config.yaml act (max_ui_session_budget: null).""" + from litellm.proxy.proxy_server import _validate_general_settings_ui_litellm_value + + assert _validate_general_settings_ui_litellm_value("max_ui_session_budget", empty_value) == 1.0 + + +def test_general_settings_ui_defaults_unchanged_for_existing_fields(): + """The spec-default mechanism added for max_ui_session_budget must not change what + clearing the pre-existing fields restores (None for Float/Select, False for Boolean).""" + from litellm.proxy.proxy_server import ( + _GENERAL_SETTINGS_UI_LITELLM_FIELDS, + _general_settings_ui_litellm_default, + ) + + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) is None + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) is False + assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) is None + + @pytest.mark.parametrize( "field_name, db_value", [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index dda7a23a8d4..6f05c896a1f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -75,6 +75,17 @@ const SettingValueEditor: React.FC<{ /> ); } + if (setting.field_type === "Dollar") { + return ( + onChange(setting.field_name, newValue)} + /> + ); + } if (setting.field_type === "Select") { return ( Date: Tue, 21 Jul 2026 21:04:11 -0700 Subject: [PATCH 030/130] fix(ui): resolve General settings rows by field name, not filtered index The General tab renders generalSettings with TypedDictionary and prompt-caching rows filtered out, but the Update and Reset handlers indexed into the unfiltered array, so any row rendered after a filtered-out entry read another field's value. max_ui_session_budget is the first General-tab row positioned after the prompt-caching entries, so its Update sent that row's boolean and failed Dollar validation. Reset also cleared the local input to null, which reads as unset or unlimited while the backend had restored the default. Handlers now resolve the row by field name and drop the index parameter, and reset displays the row's field_default_value. Component tests drive the real /config/list ordering through the actual clicks and fail under either original behavior. --- .../_components/general_settings.test.tsx | 101 ++++++++++++++++++ .../_components/general_settings.tsx | 15 +-- 2 files changed, 110 insertions(+), 6 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx new file mode 100644 index 00000000000..c4cfa98b2a1 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx @@ -0,0 +1,101 @@ +import { renderWithProviders, screen, within } from "../../../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { vi } from "vitest"; +import GeneralSettings from "./general_settings"; +import { deleteConfigFieldSetting, getGeneralSettingsCall, updateConfigFieldSetting } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + getGeneralSettingsCall: vi.fn(), + updateConfigFieldSetting: vi.fn().mockResolvedValue({}), + deleteConfigFieldSetting: vi.fn().mockResolvedValue({}), +})); + +vi.mock("@/components/router_settings", () => ({ default: () => null })); +vi.mock("@/components/Settings/RouterSettings/Fallbacks/Fallbacks", () => ({ default: () => null })); +vi.mock("@/components/routing_groups", () => ({ default: () => null })); + +// Mirrors the /config/list ordering: the two prompt-caching rows sit between the +// General-tab rows in the unfiltered response but are filtered out of the General +// tab's table, so any index-based lookup into the unfiltered array reads the wrong +// row for every field rendered after them. +const SETTINGS_FIXTURE = [ + { + field_name: "budget_exceeded_throttle_percentage", + field_type: "Float", + field_value: null, + field_description: "throttle fraction", + stored_in_db: null, + field_default_value: null, + }, + { + field_name: "enable_anthropic_prompt_caching", + field_type: "Boolean", + field_value: true, + field_description: "prompt caching toggle", + stored_in_db: true, + field_tab: "prompt_caching", + field_default_value: false, + }, + { + field_name: "anthropic_prompt_caching_ttl", + field_type: "Select", + field_value: "5m", + field_description: "prompt caching ttl", + stored_in_db: true, + field_options: ["5m", "1h"], + field_tab: "prompt_caching", + field_default_value: null, + }, + { + field_name: "max_ui_session_budget", + field_type: "Dollar", + field_value: 7.5, + field_description: "dashboard session budget", + stored_in_db: true, + field_default_value: 1.0, + }, +]; + +const settingsRow = async (fieldName: string) => { + const cell = await screen.findByText(fieldName); + const row = cell.closest("tr"); + expect(row).not.toBeNull(); + return row as HTMLElement; +}; + +describe("GeneralSettings General tab", () => { + beforeEach(() => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([...SETTINGS_FIXTURE.map((s) => ({ ...s }))]); + vi.mocked(updateConfigFieldSetting).mockClear(); + vi.mocked(deleteConfigFieldSetting).mockClear(); + }); + + it("updates max_ui_session_budget with its own value, not the value at its filtered index", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("General")); + const row = await settingsRow("max_ui_session_budget"); + + await user.click(within(row).getByRole("button", { name: /update/i })); + + expect(updateConfigFieldSetting).toHaveBeenCalledWith("token", "max_ui_session_budget", 7.5); + }); + + it("reset shows the field's default value instead of an empty input", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("General")); + const row = await settingsRow("max_ui_session_budget"); + expect(within(row).getByRole("spinbutton")).toHaveValue("7.50"); + + const actionCell = row.querySelectorAll("td")[3]; + const resetIcon = actionCell.querySelector("svg"); + expect(resetIcon).not.toBeNull(); + await user.click(resetIcon as unknown as Element); + + expect(deleteConfigFieldSetting).toHaveBeenCalledWith("token", "max_ui_session_budget"); + expect(within(row).getByRole("spinbutton")).toHaveValue("1.00"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 6f05c896a1f..fa3447e0cbf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -41,6 +41,7 @@ export interface generalSettingsItem { stored_in_db: boolean | null; field_options?: string[] | null; field_tab?: string | null; + field_default_value?: any; } const SettingValueEditor: React.FC<{ @@ -182,12 +183,12 @@ const GeneralSettings: React.FC = ({ accessToken, user setGeneralSettings(updatedSettings); }; - const handleUpdateField = (fieldName: string, idx: number) => { + const handleUpdateField = (fieldName: string) => { if (!accessToken) { return; } - let fieldValue = generalSettings[idx].field_value; + let fieldValue = generalSettings.find((setting) => setting.field_name === fieldName)?.field_value; if (fieldValue == null || fieldValue == undefined) { return; @@ -205,7 +206,7 @@ const GeneralSettings: React.FC = ({ accessToken, user } }; - const handleResetField = (fieldName: string, idx: number) => { + const handleResetField = (fieldName: string) => { if (!accessToken) { return; } @@ -215,7 +216,9 @@ const GeneralSettings: React.FC = ({ accessToken, user // update value in state const updatedSettings = generalSettings.map((setting) => - setting.field_name === fieldName ? { ...setting, stored_in_db: null, field_value: null } : setting, + setting.field_name === fieldName + ? { ...setting, stored_in_db: null, field_value: setting.field_default_value ?? null } + : setting, ); setGeneralSettings(updatedSettings); } catch (error) { @@ -292,8 +295,8 @@ const GeneralSettings: React.FC = ({ accessToken, user )} - - handleResetField(value.field_name, index)}> + + handleResetField(value.field_name)}> Reset From cfbcef319c61ad29fbb2caf9a16ac5419006bb2b Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 21 Jul 2026 22:24:14 -0700 Subject: [PATCH 031/130] fix(mcp): log actionable OAuth discovery failures for misconfigured server urls A typo'd MCP server url failed OAuth endpoint discovery silently: every failure died at debug level, the config loader warned nothing, and the /authorize 400 blamed "servers with no url" even when a url was set. _descovery_metadata now records each attempt's outcome and, when a total failure would leave the server's flow without a needed endpoint, logs one warning with the trail (urls origin-only, exception text url-stripped). Both server loaders warn which endpoints stayed unresolved for the server's flow (client_credentials never needs authorization_url, OBO needs only token_url) with the remedies; this replaces the DB path's reason-less warning and closes the config path's no-warning gap. The authorize/token/register 400 details branch on server shape via one shared helper and point at the proxy logs. _redact_mcp_resource_url moves to oauth_utils.py so the manager can import it without a cycle. Resolves LIT-4658 --- .../mcp_server/discoverable_endpoints.py | 56 ++- .../mcp_server/mcp_server_manager.py | 355 ++++++++++++++---- .../_experimental/mcp_server/oauth_utils.py | 23 +- .../proxy/_experimental/mcp_server/server.py | 25 +- .../mcp_server/test_discoverable_endpoints.py | 114 ++++++ .../mcp_server/test_mcp_server_manager.py | 175 ++++++++- 6 files changed, 623 insertions(+), 125 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 9a1b5cf4864..bc0822a71d4 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -499,6 +499,35 @@ def _raise_if_not_oauth2(mcp_server: MCPServer) -> None: ) +def _endpoint_not_configured_detail( + mcp_server: MCPServer, + endpoint_label: str, + manual_remedy: str, + issuer_remedy: str, +) -> str: + """The 400 detail for an unresolved OAuth endpoint, naming the likely cause for this server's + shape (LIT-4658): an anchored issuer whose metadata fell short, a configured (possibly + misconfigured) server url whose discovery failed, or no discovery source at all. Kept free of + URLs and issuer values because these endpoints are reachable pre-auth.""" + if mcp_server.issuer_is_anchored: + return ( + f"MCP server {endpoint_label} is not configured. Endpoint discovery anchored on the configured " + f"Issuer (RFC 8414) failed or its metadata did not include this endpoint; check the proxy logs " + f"for 'MCP OAuth' warnings from server load, verify the Issuer, or {manual_remedy}." + ) + if mcp_server.url: + return ( + f"MCP server {endpoint_label} is not configured. OAuth endpoint discovery against the configured " + f"server url did not resolve it; the url may be misconfigured. Check the proxy logs for " + f"'MCP OAuth' warnings from server load, verify the server url, or {manual_remedy}, or " + f"{issuer_remedy}." + ) + return ( + f"MCP server {endpoint_label} is not configured. Servers with no url (OpenAPI spec or stdio) run no " + f"resource discovery, so {manual_remedy}, or {issuer_remedy}." + ) + + def _raise_unless_oauth2_discovery_server( mcp_server: Optional[MCPServer], mcp_server_name: Optional[str], @@ -599,10 +628,11 @@ async def authorize_with_server( if mcp_server.authorization_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server authorization url is not configured. Servers with no url (OpenAPI " - "spec or stdio) run no resource discovery, so set Authorization URL and Token URL " - "manually, or set Issuer to discover them from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "authorization url", + "set Authorization URL and Token URL manually", + "set Issuer to discover them from the identity provider (RFC 8414)", ), ) @@ -711,10 +741,11 @@ async def exchange_token_with_server( if mcp_server.token_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server token url is not configured. Servers with no url (OpenAPI spec or " - "stdio) run no resource discovery, so set Token URL manually, or set Issuer to " - "discover it from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "token url", + "set Token URL manually", + "set Issuer to discover it from the identity provider (RFC 8414)", ), ) @@ -1278,10 +1309,11 @@ async def register_client_with_server( if mcp_server.authorization_url is None: raise HTTPException( status_code=400, - detail=( - "MCP server authorization url is not configured. Servers with no url (OpenAPI " - "spec or stdio) run no resource discovery, so set Authorization URL and Token URL " - "manually, or set Issuer to discover them from the identity provider (RFC 8414)." + detail=_endpoint_not_configured_detail( + mcp_server, + "authorization url", + "set Authorization URL and Token URL manually", + "set Issuer to discover them from the identity provider (RFC 8414)", ), ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 90b70dd01f2..4faf6da340f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,6 +13,7 @@ import json import os import re import time +from collections.abc import Sequence from contextlib import asynccontextmanager from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast from urllib.parse import urlparse @@ -50,6 +51,9 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) +from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + MCP_ELICITATION_AVAILABLE, +) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, MCPUpstreamAuthError, @@ -59,17 +63,14 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) -from litellm.proxy._experimental.mcp_server.elicitation_handler import ( - MCP_ELICITATION_AVAILABLE, -) -from litellm.proxy._experimental.mcp_server.sampling_handler import ( - MCP_SAMPLING_AVAILABLE, -) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, resolve_mcp_auth, ) +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _redact_mcp_resource_url, +) from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, Ok, @@ -100,6 +101,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ServerSpec, TokenExchangeConfig, ) +from litellm.proxy._experimental.mcp_server.sampling_handler import ( + MCP_SAMPLING_AVAILABLE, +) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -143,11 +147,9 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes try: - from mcp.shared.tool_name_validation import ( - validate_tool_name, # pyright: ignore[reportAssignmentType] - ) from mcp.shared.tool_name_validation import ( SEP_986_URL, + validate_tool_name, # pyright: ignore[reportAssignmentType] ) except ImportError: from pydantic import BaseModel @@ -408,6 +410,88 @@ def _restrict_discovery_to_corroborated_authorization_server( return metadata.model_copy(update={"token_url": None, "registration_url": None}) +def _redacted_origin_list(urls: Sequence[str]) -> str: + return ", ".join(_redact_mcp_resource_url(url) or "" for url in urls) + + +def _sanitized_error_text(exc: Exception) -> str: + return re.sub(r"https?://\S+", "", str(exc))[:200] + + +def _discovery_failure_leaves_needs_unresolved( + *, + needs_authorization_url: bool, + needs_token_url: bool, + manual_authorization_url: str | None, + manual_token_url: str | None, +) -> bool: + return (needs_authorization_url and not manual_authorization_url) or (needs_token_url and not manual_token_url) + + +def _warn_oauth_endpoints_unresolved( + *, + server_ref: str, + server_url: str | None, + discovery_attempted: bool, + issuer_anchored: bool, + metadata: MCPOAuthMetadata | None, + needs_authorization_url: bool, + needs_token_url: bool, + manual_authorization_url: str | None, + manual_token_url: str | None, +) -> None: + """Log one actionable warning when a server that depends on OAuth endpoint discovery finishes a + build without the endpoints that its flows need (LIT-4658). + + This is the operator-facing signal for a misconfigured server url: discovery failures themselves + are logged where they happen (``_descovery_metadata``), and this names WHICH server is affected, + which endpoints stayed unresolved after manual configuration was considered, and the remedies. + Scopes never trigger the warning on their own: scope-less metadata is normal for many servers and + warning on it every rebuild would be noise. Callers own the per-flow policy of which endpoints + are needed (client_credentials never needs authorization_url; OBO needs only token_url); the + issuer-anchored arm is excluded here because it has its own RFC 8414 §3.3 warning. + """ + if issuer_anchored: + return + unresolved = tuple( + field + for field, needed, value in ( + ( + "authorization_url", + needs_authorization_url, + manual_authorization_url or (metadata.authorization_url if metadata else None), + ), + ( + "token_url", + needs_token_url, + manual_token_url or (metadata.token_url if metadata else None), + ), + ) + if needed and not value + ) + if not unresolved: + return + if discovery_attempted: + verbose_logger.warning( + "MCP server %s: OAuth endpoint discovery left %s unresolved (server url origin: %s). OAuth flows " + "that need them will fail with 'not configured' errors until they resolve. Check the preceding " + "'MCP OAuth' log lines for why discovery failed, verify the configured server url, or set the " + "unresolved endpoint urls manually, or set issuer to discover them from the identity provider " + "(RFC 8414)", + server_ref, + ", ".join(unresolved), + _redact_mcp_resource_url(server_url) or "", + ) + return + verbose_logger.warning( + "MCP server %s uses OAuth but has no discovery source (no server url or pinned issuer), and %s not " + "set manually. Set the missing endpoint urls on the server, or set issuer to discover them from the " + "identity provider (RFC 8414)", + server_ref, + " and ".join(unresolved) + (" is" if len(unresolved) == 1 else " are"), + ) + + def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: """Drop a cached entry after the user stores or clears their env var values so the next request reads the fresh value instead of a stale one.""" @@ -884,10 +968,10 @@ def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): return None async def _sampling_callback(context, params): + import litellm from litellm.proxy._experimental.mcp_server.sampling_handler import ( handle_sampling_create_message, ) - import litellm from litellm.proxy._experimental.mcp_server.server import ( get_active_auth_context, ) @@ -1284,6 +1368,15 @@ class MCPServerManager: should_discover = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( is_discovery_auth_type or obo_needs_discovery ) + config_oauth2_flow = server_config.get("oauth2_flow", None) + needs_authorization_url = is_discovery_auth_type and config_oauth2_flow != "client_credentials" + needs_token_url = is_discovery_auth_type or obo_needs_discovery + warn_on_empty_discovery = _discovery_failure_leaves_needs_unresolved( + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) if not should_discover: mcp_oauth_metadata = None elif use_issuer_anchor and manual_issuer is not None: @@ -1292,6 +1385,7 @@ class MCPServerManager: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, allow_origin_fallback=is_discovery_auth_type, + warn_when_no_metadata=warn_on_empty_discovery, ) if use_issuer_anchor: @@ -1326,7 +1420,6 @@ class MCPServerManager: ) effective_issuer = manual_issuer or discovered_issuer - config_oauth2_flow = server_config.get("oauth2_flow", None) if auth_type == MCPAuth.oauth2 and config_oauth2_flow not in ( "client_credentials", "authorization_code", @@ -1358,6 +1451,18 @@ class MCPServerManager: "authorization-code flow." ) + _warn_oauth_endpoints_unresolved( + server_ref=server_name or server_id, + server_url=server_url, + discovery_attempted=should_discover, + issuer_anchored=use_issuer_anchor, + metadata=gated_oauth_metadata, + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + new_server = MCPServer( server_id=server_id, name=name_for_prefix, @@ -1485,14 +1590,12 @@ class MCPServerManager: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( build_input_schema, create_tool_function, + load_openapi_spec_async, + resolve_operation_params, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( get_base_url as get_openapi_base_url, ) - from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - load_openapi_spec_async, - resolve_operation_params, - ) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -1681,10 +1784,20 @@ class MCPServerManager: scopes: Optional[list[str]], token_exchange_endpoint: Optional[str], ) -> Optional[MCPOAuthMetadata]: + obo_needs_discovery = self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) + needs_authorization_url = ( + is_discovery_auth_type and getattr(mcp_server, "oauth2_flow", None) != "client_credentials" + ) + needs_token_url = is_discovery_auth_type or obo_needs_discovery + warn_on_empty_discovery = _discovery_failure_leaves_needs_unresolved( + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) needs_discovery = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( - (is_discovery_auth_type and not has_all_upstream_oauth_fields) - or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) + (is_discovery_auth_type and not has_all_upstream_oauth_fields) or obo_needs_discovery ) if not needs_discovery: mcp_oauth_metadata: Optional[MCPOAuthMetadata] = None @@ -1694,24 +1807,32 @@ class MCPServerManager: mcp_oauth_metadata = await self._descovery_metadata( server_url=server_url, # type: ignore[arg-type] allow_origin_fallback=is_discovery_auth_type, - ) - if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None: - verbose_logger.warning( - "MCP OAuth discovery yielded no metadata for server %s (%s); " - "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", - mcp_server.server_id, - server_url, + warn_when_no_metadata=warn_on_empty_discovery, ) if use_issuer_anchor: return mcp_oauth_metadata - if is_discovery_auth_type: - return _restrict_discovery_to_corroborated_authorization_server( + gated_metadata = ( + _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, manual_authorization_url, mcp_server.server_id, bool(getattr(mcp_server, "dcr_bridge", None)), ) - return mcp_oauth_metadata + if is_discovery_auth_type + else mcp_oauth_metadata + ) + _warn_oauth_endpoints_unresolved( + server_ref=mcp_server.alias or mcp_server.server_name or mcp_server.server_id, + server_url=server_url, + discovery_attempted=needs_discovery, + issuer_anchored=False, + metadata=gated_metadata, + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + return gated_metadata async def build_mcp_server_from_table( self, @@ -3430,6 +3551,7 @@ class MCPServerManager: server_url: str, *, allow_origin_fallback: bool = True, + warn_when_no_metadata: bool = False, ) -> Optional[MCPOAuthMetadata]: """Discover OAuth metadata by following RFC 9728 (protected resource metadata discovery). @@ -3438,8 +3560,32 @@ class MCPServerManager: it (a human sees the redirect), but token_exchange (OBO) sets it False so the gateway never exchanges a subject token against an endpoint it inferred rather than one explicitly configured or authoritatively advertised via RFC 9728 / RFC 8414. - """ + ``warn_when_no_metadata`` makes an all-empty result log one WARNING with the per-step attempt + outcomes (LIT-4658), so a misconfigured server url is diagnosable from default-level logs. The + server loaders set it; the issuer-anchored resource-scopes lookup keeps it off because empty + scopes are not a fault there. + """ + metadata, attempts = await self._discover_metadata_recording_attempts( + server_url, allow_origin_fallback=allow_origin_fallback + ) + if metadata is None and warn_when_no_metadata: + verbose_logger.warning( + "MCP OAuth endpoint discovery against %s found no authorization server metadata. Attempts: %s. " + "The MCP server url may be misconfigured, or the upstream may not support OAuth discovery " + "(RFC 9728 / RFC 8414)", + _redact_mcp_resource_url(server_url) or "", + "; ".join(attempts) if attempts else "none recorded", + ) + return metadata + + async def _discover_metadata_recording_attempts( + self, + server_url: str, + *, + allow_origin_fallback: bool, + ) -> tuple[MCPOAuthMetadata | None, tuple[str, ...]]: + origin = _redact_mcp_resource_url(server_url) or "" try: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) response = await client.get(server_url) @@ -3452,67 +3598,112 @@ class MCPServerManager: if metadata is None and not resource_scopes and authorization_servers and response.status_code == 200: verbose_logger.warning( "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", - server_url, + origin, ) + attempts = ( + f"GET {origin}: HTTP {response.status_code} (no RFC 9728 challenge)", + *( + ("well-known protected-resource lookup found no authorization servers",) + if not authorization_servers + else () + ), + *( + (f"authorization server metadata fetch failed for: {_redacted_origin_list(authorization_servers)}",) + if authorization_servers and metadata is None + else () + ), + ) if metadata is None and resource_scopes: - return MCPOAuthMetadata(scopes=resource_scopes) + return MCPOAuthMetadata(scopes=resource_scopes), attempts if metadata is not None and resource_scopes: metadata.scopes = resource_scopes - return metadata + return metadata, attempts except HTTPStatusError as exc: - verbose_logger.debug( - "MCP OAuth discovery for %s received status error: %s", - server_url, - exc, - ) - - header_value: Optional[str] = None - if exc.response is not None: - header_value = exc.response.headers.get("WWW-Authenticate") or exc.response.headers.get( - "www-authenticate" - ) - - resource_metadata_url, scopes = self._parse_www_authenticate_header(header_value) - - authorization_servers = [] - resource_scopes = None - if resource_metadata_url: - ( - authorization_servers, - resource_scopes, - ) = await self._fetch_oauth_metadata_from_resource(resource_metadata_url, server_url) - else: - ( - authorization_servers, - resource_scopes, - ) = await self._attempt_well_known_discovery(server_url) - - metadata = None - used_origin_fallback = False - if allow_origin_fallback and not authorization_servers: - try: - parsed_url = urlparse(server_url) - if parsed_url.scheme and parsed_url.netloc: - authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"] - used_origin_fallback = True - except Exception: - authorization_servers = [] - - if authorization_servers: - metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) - if metadata is not None and used_origin_fallback: - metadata.from_origin_fallback = True - - preferred_scopes = scopes or resource_scopes - if metadata is None and preferred_scopes: - metadata = MCPOAuthMetadata(scopes=preferred_scopes) - elif metadata is not None and preferred_scopes: - metadata.scopes = preferred_scopes - - return metadata + return await self._discover_after_status_error(server_url, exc, allow_origin_fallback=allow_origin_fallback) except Exception as exc: # pragma: no cover - network/transient issues verbose_logger.debug("MCP OAuth discovery failed for %s: %s", server_url, exc) - return None + return None, (f"GET {origin}: {type(exc).__name__}: {_sanitized_error_text(exc)}",) + + async def _discover_after_status_error( + self, + server_url: str, + exc: HTTPStatusError, + *, + allow_origin_fallback: bool, + ) -> tuple[MCPOAuthMetadata | None, tuple[str, ...]]: + origin = _redact_mcp_resource_url(server_url) or "" + verbose_logger.debug( + "MCP OAuth discovery for %s received status error: %s", + server_url, + exc, + ) + + header_value: Optional[str] = None + if exc.response is not None: + header_value = exc.response.headers.get("WWW-Authenticate") or exc.response.headers.get("www-authenticate") + status_attempt = ( + f"GET {origin}: HTTP {exc.response.status_code}" + if exc.response is not None + else f"GET {origin}: status error" + ) + + resource_metadata_url, scopes = self._parse_www_authenticate_header(header_value) + + authorization_servers = [] + resource_scopes = None + if resource_metadata_url: + ( + authorization_servers, + resource_scopes, + ) = await self._fetch_oauth_metadata_from_resource(resource_metadata_url, server_url) + lookup_attempt = ( + None + if authorization_servers + else "challenge-advertised resource metadata yielded no authorization servers" + ) + else: + ( + authorization_servers, + resource_scopes, + ) = await self._attempt_well_known_discovery(server_url) + lookup_attempt = ( + None + if authorization_servers + else "no challenge-advertised resource metadata; well-known protected-resource lookup found no authorization servers" + ) + + metadata = None + used_origin_fallback = False + if allow_origin_fallback and not authorization_servers: + try: + parsed_url = urlparse(server_url) + if parsed_url.scheme and parsed_url.netloc: + authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"] + used_origin_fallback = True + except Exception: + authorization_servers = [] + + fallback_attempt = None + if authorization_servers: + metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) + if metadata is not None and used_origin_fallback: + metadata.from_origin_fallback = True + if metadata is None: + fallback_attempt = ( + f"origin fallback: no authorization server metadata at {origin}" + if used_origin_fallback + else f"authorization server metadata fetch failed for: {_redacted_origin_list(authorization_servers)}" + ) + + attempts = tuple(entry for entry in (status_attempt, lookup_attempt, fallback_attempt) if entry) + + preferred_scopes = scopes or resource_scopes + if metadata is None and preferred_scopes: + return MCPOAuthMetadata(scopes=preferred_scopes), attempts + if metadata is not None and preferred_scopes: + metadata.scopes = preferred_scopes + + return metadata, attempts def _parse_www_authenticate_header(self, header_value: Optional[str]) -> tuple[Optional[str], Optional[list[str]]]: if not header_value: diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 53686e329bb..2f92f75a352 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -4,7 +4,7 @@ import os from ipaddress import ip_address from typing import Any, Dict, List, NoReturn, Optional -from urllib.parse import ParseResult, urlparse, urlunparse +from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit from fastapi import HTTPException, Request @@ -70,6 +70,27 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" +def _redact_mcp_resource_url(url: Optional[str]) -> Optional[str]: + """Reduce an MCP server URL to its origin (scheme + host + port) for logging. + + Everything else is dropped: userinfo (``user:pass@``), the query string, the + fragment, and the path, because hosted MCP servers routinely embed the + credential in the path (e.g. ``/mcp/s/``) and this value is persisted + in spend-log metadata that a caller who can invoke the tool can read back. + Returns None when the URL has no host to identify (nothing safe to log). + """ + if not isinstance(url, str) or not url: + return None + try: + parts = urlsplit(url) + except ValueError: + return None + if not parts.hostname: + return None + netloc = f"{parts.hostname}:{parts.port}" if parts.port else parts.hostname + return urlunsplit((parts.scheme, netloc, "", "", "")) or None + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 396dd6c7dc7..56e8ec25076 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -27,7 +27,6 @@ from typing import ( Union, cast, ) -from urllib.parse import urlsplit, urlunsplit import httpx from fastapi import FastAPI, HTTPException @@ -59,6 +58,9 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _redact_mcp_resource_url, +) from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -106,27 +108,6 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100 _MCP_ROUTING_PEEK_MAX_BYTES = 4096 -def _redact_mcp_resource_url(url: Optional[str]) -> Optional[str]: - """Reduce an MCP server URL to its origin (scheme + host + port) for logging. - - Everything else is dropped: userinfo (``user:pass@``), the query string, the - fragment, and the path, because hosted MCP servers routinely embed the - credential in the path (e.g. ``/mcp/s/``) and this value is persisted - in spend-log metadata that a caller who can invoke the tool can read back. - Returns None when the URL has no host to identify (nothing safe to log). - """ - if not isinstance(url, str) or not url: - return None - try: - parts = urlsplit(url) - except ValueError: - return None - if not parts.hostname: - return None - netloc = f"{parts.hostname}:{parts.port}" if parts.port else parts.hostname - return urlunsplit((parts.scheme, netloc, "", "", "")) or None - - def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 636c7fbd3d5..eeeea82b647 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -8148,3 +8148,117 @@ async def test_register_wall_names_the_fix_for_urlless_servers(): detail_text = str(exc_info.value.detail) assert "set Authorization URL and Token URL" in detail_text assert "Issuer" in detail_text + + +@pytest.mark.asyncio +async def test_authorize_wall_points_at_discovery_failure_for_url_servers(): + """LIT-4658: a server WITH a url that still has no authorization_url got here because OAuth + discovery against that url failed (typically a misconfigured url); the old detail blamed + "servers with no url", sending the operator down the wrong path. The detail must now name the + discovery failure and point at the proxy logs where LIT-4658's warnings carry the reason.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="typo-url-wall", + name="typo_wall", + server_name="typo_wall", + url="https://typo-host.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="client", + redirect_uri="http://localhost/callback", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "may be misconfigured" in detail_text + assert "proxy logs" in detail_text + assert "Servers with no url" not in detail_text + assert "typo-host.example.com" not in detail_text + + +@pytest.mark.asyncio +async def test_token_wall_points_at_discovery_failure_for_url_servers(): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="typo-url-token-wall", + name="typo_token_wall", + server_name="typo_token_wall", + url="https://typo-host.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="http://localhost/callback", + client_id="client", + client_secret=None, + code_verifier="verifier", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "token url is not configured" in detail_text + assert "may be misconfigured" in detail_text + assert "Servers with no url" not in detail_text + + +@pytest.mark.asyncio +async def test_authorize_wall_names_the_issuer_for_anchored_servers(): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="anchored-wall", + name="anchored_wall", + server_name="anchored_wall", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + issuer="https://idp.example.com", + issuer_is_anchored=True, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="client", + redirect_uri="http://localhost/callback", + ) + assert exc_info.value.status_code == 400 + detail_text = str(exc_info.value.detail) + assert "verify the Issuer" in detail_text + assert "Servers with no url" not in detail_text + assert "idp.example.com" not in detail_text 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 a5cb16822cf..77f072b81d5 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 @@ -3031,7 +3031,7 @@ class TestMCPServerManager: registration_url="https://discovered.example.com/register", ) - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert server_url == "https://example.com/mcp" # oauth2 (browser flow) keeps the origin fallback; only OBO disables it. assert allow_origin_fallback is True @@ -5426,7 +5426,7 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[bool] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): calls.append(allow_origin_fallback) return MCPOAuthMetadata( scopes=None, @@ -5461,7 +5461,7 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[str] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): calls.append(server_url) raise AssertionError("discovery must not run when token_exchange_endpoint is configured") @@ -5491,7 +5491,7 @@ class TestMCPServerTimestamps: back to the row, so the next rebuild skips discovery instead of re-running it every time.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert server_url == "https://example.com/mcp" assert allow_origin_fallback is False # OBO never guesses the origin return MCPOAuthMetadata( @@ -5602,7 +5602,7 @@ class TestMCPServerTimestamps: _dcr_bridge_relays_client_registration keys off that column.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): assert allow_origin_fallback is True return MCPOAuthMetadata( scopes=["mcp.read", "mcp.write"], @@ -5817,7 +5817,7 @@ class TestMCPServerTimestamps: persist_discovered_endpoints=False neither the oauth2 nor the OBO write-back may fire.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): return MCPOAuthMetadata( scopes=["s1"], authorization_url="https://idp.example.com/authorize", @@ -8539,7 +8539,7 @@ class TestOBOEndpointDiscovery: ) seen = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): seen.append((server_url, allow_origin_fallback)) return discovered @@ -8567,7 +8567,7 @@ class TestOBOEndpointDiscovery: async def test_config_obo_with_configured_endpoint_skips_discovery(self): manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True): + async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): raise AssertionError("discovery must not run when the endpoint is configured") manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] @@ -9028,3 +9028,162 @@ class TestUrllessIssuerDiscovery: anchored.assert_awaited_once_with("https://idp.example.com", None) resource_rooted.assert_not_awaited() assert built.token_url == "https://idp.example.com/token" + + +class TestDiscoveryFailureLogging: + """LIT-4658: a misconfigured MCP server url must be diagnosable from default-level server logs. + + Discovery failures used to die at debug level and the config-load path emitted no warning at + all, so the only operator-facing signal was the bare 400 at /authorize.""" + + def _connect_error_client(self, url: str) -> MagicMock: + client = MagicMock() + client.get = AsyncMock( + side_effect=httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}") + ) + return client + + @pytest.mark.asyncio + async def test_descovery_metadata_warns_with_redacted_attempts_on_connect_error(self, caplog): + manager = MCPServerManager() + secret_url = "https://typo-host.example.com/mcp/s/PATHSECRET/mcp" + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=self._connect_error_client(secret_url), + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(secret_url, warn_when_no_metadata=True) + assert result is None + assert "found no authorization server metadata" in caplog.text + assert "ConnectError" in caplog.text + assert "https://typo-host.example.com" in caplog.text + # hosted MCP urls embed credentials in the path; neither the url nor the exception + # text may leak it into warning-level logs + assert "PATHSECRET" not in caplog.text + + @pytest.mark.asyncio + async def test_descovery_metadata_stays_silent_without_warn_flag(self, caplog): + manager = MCPServerManager() + url = "https://typo-host.example.com/mcp" + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=self._connect_error_client(url), + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(url) + assert result is None + assert "found no authorization server metadata" not in caplog.text + + @pytest.mark.asyncio + async def test_descovery_metadata_attempt_trail_names_each_failed_step(self, caplog): + manager = MCPServerManager() + url = "https://real-host.example.com/mcp-typo" + client = MagicMock() + client.get = AsyncMock( + return_value=httpx.Response(404, request=httpx.Request("GET", url)) + ) + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=client, + ), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + result = await manager._descovery_metadata(url, warn_when_no_metadata=True) + assert result is None + assert "HTTP 404" in caplog.text + assert "well-known protected-resource lookup found no authorization servers" in caplog.text + assert "origin fallback" in caplog.text + + @pytest.mark.asyncio + async def test_load_servers_from_config_warns_when_endpoints_unresolved(self, caplog): + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + config = { + "typo_server": { + "url": "https://typo.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "typo_server" in caplog.text + assert "authorization_url, token_url" in caplog.text + assert "unresolved" in caplog.text + assert "verify the configured server url" in caplog.text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "extra_config", + [ + { + "authorization_url": "https://idp.example.com/auth", + "token_url": "https://idp.example.com/token", + }, + { + "oauth2_flow": "client_credentials", + "token_url": "https://idp.example.com/token", + "client_id": "cid", + "client_secret": "csec", + }, + ], + ) + async def test_load_servers_from_config_silent_when_flow_needs_covered(self, caplog, extra_config): + """Manually covered endpoints and M2M servers (which never need authorization_url) must not + warn on every reload; the warning is a misconfiguration signal, not discovery telemetry.""" + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + config = { + "covered_server": { + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + **extra_config, + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "unresolved" not in caplog.text + assert "no discovery source" not in caplog.text + + @pytest.mark.asyncio + async def test_config_server_without_discovery_source_warns_about_missing_endpoints(self, caplog): + manager = MCPServerManager() + manager._register_openapi_tools = AsyncMock() # type: ignore[attr-defined] + config = { + "spec_only": { + "spec_path": "https://example.com/openapi.yaml", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + } + } + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + assert "no discovery source" in caplog.text + assert "authorization_url and token_url are not set manually" in caplog.text + + @pytest.mark.asyncio + async def test_db_build_warns_when_discovery_fails_for_oauth2_row(self, caplog): + manager = MCPServerManager() + manager._descovery_metadata = AsyncMock(return_value=None) # type: ignore[attr-defined] + record = LiteLLM_MCPServerTable( + server_id="typo-row-1", + server_name="typo_row", + url="https://typo.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + assert "typo_row" in caplog.text + assert "authorization_url, token_url" in caplog.text + assert "unresolved" in caplog.text From df75d298ec25db8cb69560b7ca3c266ba84636ea Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 21 Jul 2026 22:45:39 -0700 Subject: [PATCH 032/130] fix(mcp): keep url redaction total when the port is malformed urlsplit validates the port lazily, so a non-numeric port raised ValueError out of _redact_mcp_resource_url after the urlsplit try had already passed; the server loaders now call the helper while warning about typo'd urls, which would have turned the warning into a load failure. Resolve hostname and port inside the guard and pin the malformed-port case in the redaction test --- litellm/proxy/_experimental/mcp_server/oauth_utils.py | 6 ++++-- .../proxy/_experimental/mcp_server/test_mcp_server.py | 4 ++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 2f92f75a352..8a5f398003b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -83,11 +83,13 @@ def _redact_mcp_resource_url(url: Optional[str]) -> Optional[str]: return None try: parts = urlsplit(url) + hostname = parts.hostname + port = parts.port except ValueError: return None - if not parts.hostname: + if not hostname: return None - netloc = f"{parts.hostname}:{parts.port}" if parts.port else parts.hostname + netloc = f"{hostname}:{port}" if port else hostname return urlunsplit((parts.scheme, netloc, "", "", "")) or None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ae4f12fc1e1..dff1f1d87c7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7375,6 +7375,10 @@ async def test_call_tool_with_legacy_db_m2m_server_resolves_oauth2_flow(): ("", None), ("not a url", None), ("http://[::1", None), + # urlsplit validates the port lazily on attribute access, so a malformed port must not + # raise out of the helper: the server loaders call it while warning about exactly this + # kind of typo'd url (LIT-4658) + ("https://example.com:bad/mcp", None), ], ) def test_redact_mcp_resource_url_strips_credentials(url, expected): From f1f0a0bacd5ce681d404264db2bcd27801d9d8ac Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 21 Jul 2026 23:15:37 -0700 Subject: [PATCH 033/130] fix(bedrock): emit Nova Sonic realtime session.created on connect and session.updated on session.update (#34133) --- litellm/llms/bedrock/realtime/handler.py | 26 ++ .../llms/bedrock/realtime/transformation.py | 45 ++-- pyproject.toml | 8 + .../llm_nonconversational.yaml | 2 +- ...me_e2e.py => test_realtime_bedrock_e2e.py} | 0 .../realtime/test_bedrock_realtime_handler.py | 72 +++++- .../test_bedrock_realtime_transformation.py | 53 +++- uv.lock | 229 +++++++++++++++++- 8 files changed, 405 insertions(+), 30 deletions(-) rename tests/e2e/llm_translation/realtime/{test_nova_sonic_realtime_e2e.py => test_realtime_bedrock_e2e.py} (100%) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index b48c37791c4..b7237d288ec 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -9,6 +9,8 @@ import contextlib import json from typing import Any, Optional +from pydantic import TypeAdapter + from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -16,6 +18,8 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +_CLIENT_MODALITIES_ADAPTER: TypeAdapter["list[str] | None"] = TypeAdapter(list[str] | None) + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -124,6 +128,9 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) + verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") + # Track state for transformation session_state = { "current_output_item_id": None, @@ -143,6 +150,7 @@ class BedrockRealtime(BaseAWSLLM): transformation_config, model, session_state, + logging_obj, ) ) @@ -179,6 +187,7 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: BedrockRealtimeConfig, model: str, session_state: dict, + logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" from aws_sdk_bedrock_runtime.models import ( @@ -210,6 +219,23 @@ class BedrockRealtime(BaseAWSLLM): for bedrock_message in transformed_messages: await send_to_bedrock(bedrock_message) + if logging_obj is not None: + client_message_type: str | None = None + requested_modalities: list[str] | None = None + with contextlib.suppress(Exception): + parsed_client_message = json.loads(message) + client_message_type = parsed_client_message.get("type") + if client_message_type == "session.update": + requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( + parsed_client_message.get("session", {}).get("modalities") + ) + if client_message_type == "session.update": + await client_ws.send_text( + json.dumps( + transformation_config.session_updated_event(model, logging_obj, requested_modalities) + ) + ) + except Exception as e: verbose_proxy_logger.debug(f"Client to Bedrock forwarding ended: {e}", exc_info=True) for close_message in transformation_config.session_close_messages(): diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index fe5f0584e03..24a40ebea1b 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -623,35 +623,42 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.warning(f"Unknown message type: {message_type}") return [] - def transform_session_start_event( + def _session_object( self, - event: dict, model: str, logging_obj: LiteLLMLoggingObj, - ) -> OpenAIRealtimeStreamSessionEvents: - """ - Transform Bedrock sessionStart event to OpenAI session.created. - - Args: - event: Bedrock sessionStart event - model: Model ID - logging_obj: Logging object - - Returns: - OpenAI session.created event - """ - verbose_logger.debug("Handling sessionStart") - + modalities: list[str] | None = None, + ) -> OpenAIRealtimeStreamSession: session = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=["text", "audio"], + modalities=modalities if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model + return session + def session_created_event( + self, + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> OpenAIRealtimeStreamSessionEvents: + """Build the OpenAI session.created event for this realtime session.""" return OpenAIRealtimeStreamSessionEvents( type="session.created", - session=session, + session=self._session_object(model, logging_obj), + event_id=str(uuid.uuid4()), + ) + + def session_updated_event( + self, + model: str, + logging_obj: LiteLLMLoggingObj, + modalities: list[str] | None = None, + ) -> OpenAIRealtimeStreamSessionEvents: + """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" + return OpenAIRealtimeStreamSessionEvents( + type="session.updated", + session=self._session_object(model, logging_obj, modalities), event_id=str(uuid.uuid4()), ) @@ -1169,8 +1176,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Route to appropriate transformation method if "sessionStart" in event: - session_created = self.transform_session_start_event(event, model, logging_obj) - returned_messages.append(session_created) session_configuration_request = json.dumps({"configured": True}) elif "contentStart" in event: diff --git a/pyproject.toml b/pyproject.toml index 080d06258ed..62bd37c3db6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,6 +117,14 @@ stt-nvidia-riva = [ "numpy>=1.26.0", ] google = ["google-cloud-aiplatform>=1.133.0,<2.0"] +bedrock-realtime = [ + # Bedrock Nova Sonic realtime (speech-to-speech) uses the + # InvokeModelWithBidirectionalStream API, which boto3 cannot do. This + # experimental AWS SDK (with its smithy-* deps, pulled transitively) + # provides the bidirectional stream; imported lazily in the realtime + # handler so litellm core stays usable without it. + "aws-sdk-bedrock-runtime>=0.7.0,<0.8.0; python_version >= '3.12'", +] proxy-runtime = [ # Historically bundled in the proxy Docker images via requirements.txt. # Keep these in a dedicated extra so uv-based images preserve the same diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 2b456aacefc..63e6fde14a3 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -32,7 +32,7 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} -- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_nova_sonic_realtime_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} +- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.rerank.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/rerank/handler.py", rationale: "Bedrock rerank"} - {id: llm.rerank.together_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: together_ai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/together_ai/rerank/handler.py", rationale: "Together rerank"} - {id: llm.images_generations.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_image_generation_e2e.py:22", rationale: "OpenAI image gen, b64/url"} diff --git a/tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_bedrock_e2e.py similarity index 100% rename from tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py rename to tests/e2e/llm_translation/realtime/test_realtime_bedrock_e2e.py diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index ddc2e026e83..ffe21b91ab2 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -54,15 +54,24 @@ class FakeBedrockStream: self.input_stream = input_stream if input_stream is not None else FakeInputStream() +class FakeLogging: + def __init__(self, trace_id="trace-nova-sonic"): + self.litellm_trace_id = trace_id + + class DisconnectingClientWS: def __init__(self, messages): self._messages = list(messages) + self.sent_to_client = [] async def receive_text(self): if self._messages: return self._messages.pop(0) raise RuntimeError("client disconnected") + async def send_text(self, message): + self.sent_to_client.append(message) + class ClosableClientWS: def __init__(self): @@ -85,10 +94,14 @@ class EndedBedrockStream: class RealtimeClientWS: def __init__(self): self.closed = False + self.sent_to_client = [] async def receive_text(self): raise RuntimeError("client disconnected") + async def send_text(self, message): + self.sent_to_client.append(message) + async def close(self, code=None, reason=None): self.closed = True @@ -277,6 +290,61 @@ class TestBedrockRealtimeHandler: assert client_ws.closed +class TestBedrockRealtimeSessionLifecycle: + """Server must emit session.created on connect and session.updated on session.update (LIT-4655 regression)""" + + @pytest.mark.asyncio + async def test_session_created_sent_on_connect_before_any_client_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = RealtimeClientWS() + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=FakeLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + + assert websocket.sent_to_client, "server sent nothing on connect: spec-conformant clients deadlock" + first_event = json.loads(websocket.sent_to_client[0]) + assert first_event["type"] == "session.created" + assert first_event["session"]["id"] == "trace-nova-sonic" + assert first_event["session"]["model"] == "amazon.nova-sonic-v1:0" + + @pytest.mark.asyncio + async def test_session_update_is_acked_with_session_updated(self, stub_aws_models): + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})] + ) + + await handler._forward_client_to_bedrock( + client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging() + ) + + acked = [json.loads(message) for message in client_ws.sent_to_client] + updated = [event for event in acked if event["type"] == "session.updated"] + assert updated, "session.update was not acked" + assert updated[0]["session"]["modalities"] == ["text"], "ack must reflect the requested modalities" + + @pytest.mark.asyncio + async def test_no_session_updated_without_logging_obj(self, stub_aws_models): + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "hi"}})] + ) + + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + + assert client_ws.sent_to_client == [] + + class TestBedrockRealtimeAwsAuth: """AWS auth params passed via litellm_params must reach the Smithy client config (LIT-3923 regression)""" @@ -288,7 +356,7 @@ class TestBedrockRealtimeAwsAuth: await handler.async_realtime( model="amazon.nova-sonic-v1:0", websocket=websocket, - logging_obj=MagicMock(), + logging_obj=FakeLogging(), aws_region_name="us-east-1", aws_access_key_id="litellm-params-access-key", aws_secret_access_key="litellm-params-secret-key", @@ -318,7 +386,7 @@ class TestBedrockRealtimeAwsAuth: await handler.async_realtime( model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), - logging_obj=MagicMock(), + logging_obj=FakeLogging(), aws_region_name="eu-west-1", aws_role_name="arn:aws:iam::123456789012:role/nova-sonic", aws_session_name="realtime-session", diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index a68aa603b26..aa002b6e302 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -403,8 +403,9 @@ class TestBedrockRealtimeResponseCreate: class TestBedrockRealtimeResponseTransformation: """Test suite for response transformation""" - def test_transform_session_start_response(self): - """Test sessionStart response transformation""" + def test_bedrock_session_start_does_not_emit_duplicate_session_created(self): + """A Bedrock output sessionStart must not forward a second session.created to the + client; session.created is sent exactly once on connect (LIT-4655)""" config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -428,10 +429,8 @@ class TestBedrockRealtimeResponseTransformation: }, ) - assert len(result["response"]) == 1 - assert result["response"][0]["type"] == "session.created" - assert result["response"][0]["session"]["id"] == "trace_123" - assert "model" in result["response"][0]["session"] + assert result["response"] == [] + assert result["session_configuration_request"] == json.dumps({"configured": True}) def test_transform_text_output_response(self): """Test textOutput response transformation""" @@ -789,5 +788,47 @@ class TestBedrockRealtimeResponseTransformation: assert len(set(response_ids)) == 1, "Response IDs should be consistent" +class TestBedrockRealtimeSessionEvents: + """session.created / session.updated builders produce spec-shaped events (LIT-4655)""" + + @staticmethod + def _logging(): + from types import SimpleNamespace + + return SimpleNamespace(litellm_trace_id="trace_123") + + def test_session_created_event_shape(self): + event = BedrockRealtimeConfig().session_created_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["type"] == "session.created" + assert event["session"]["id"] == "trace_123" + assert event["session"]["model"] == "amazon.nova-sonic-v1:0" + assert event["session"]["modalities"] == ["text", "audio"] + assert event["event_id"] + + def test_session_updated_event_shape(self): + event = BedrockRealtimeConfig().session_updated_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["type"] == "session.updated" + assert event["session"]["id"] == "trace_123" + assert event["session"]["model"] == "amazon.nova-sonic-v1:0" + assert event["event_id"] + + def test_created_and_updated_have_distinct_event_ids(self): + config = BedrockRealtimeConfig() + logging_obj = self._logging() + created = config.session_created_event("amazon.nova-sonic-v1:0", logging_obj) + updated = config.session_updated_event("amazon.nova-sonic-v1:0", logging_obj) + assert created["event_id"] != updated["event_id"] + + def test_session_updated_reflects_requested_modalities(self): + event = BedrockRealtimeConfig().session_updated_event( + "amazon.nova-sonic-v1:0", self._logging(), modalities=["text"] + ) + assert event["session"]["modalities"] == ["text"] + + def test_session_updated_defaults_modalities_when_unspecified(self): + event = BedrockRealtimeConfig().session_updated_event("amazon.nova-sonic-v1:0", self._logging()) + assert event["session"]["modalities"] == ["text", "audio"] + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/uv.lock b/uv.lock index 6de7bf20bd6..b3d6fccff26 100644 --- a/uv.lock +++ b/uv.lock @@ -530,6 +530,70 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/1f/aa74b23b6eea4cf9b79ace914df59123c4c8e7e4bd32dd22d09c126422d9/aurelio_sdk-0.0.19-py3-none-any.whl", hash = "sha256:390c0212b59ce99116df8722d3badced88c5ef0bb742a6222d479ceed0ed3948", size = 17322, upload-time = "2025-03-24T14:37:31.305Z" }, ] +[[package]] +name = "aws-sdk-bedrock-runtime" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smithy-aws-core", extra = ["eventstream", "json"], marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, + { name = "smithy-http", extra = ["awscrt"], marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/8a/ed3fd98775273b0b7f6006b4970aa876d506668b7fe29145f54fcb941c3b/aws_sdk_bedrock_runtime-0.7.0.tar.gz", hash = "sha256:0cb172cbc03ff060e5c1d6f9cfa9a8ac5e71d9e0d58d3117006ebf614cbb4677", size = 170304, upload-time = "2026-06-23T04:04:52.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/e1/f86d50f0ad9c8200645f315c524d285e86b30b94bb65118e1108597714e6/aws_sdk_bedrock_runtime-0.7.0-py3-none-any.whl", hash = "sha256:de67ede6f441bbb77ef61c237945d559513843fc827abe1af12535c2519650c5", size = 94948, upload-time = "2026-06-23T04:04:51.281Z" }, +] + +[[package]] +name = "aws-sdk-signers" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/48/62/30685f9a096ad834409771a567065cd25094d3a5124874e1d7e373f17dd7/aws_sdk_signers-0.3.0.tar.gz", hash = "sha256:6fd654ea3dafe3ae3cf172fb3bd1e81877fe6e18156e0b638ccefb045ffb420c", size = 18590, upload-time = "2026-05-05T18:04:09.914Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/30/91/25405e647de8dea8419812e85bcb341323500905f4f55fc23cf488b7a2a4/aws_sdk_signers-0.3.0-py3-none-any.whl", hash = "sha256:098c7784064931d2da968400fc0466e092a7c271a79fac9d55daa4a7c35fd6f1", size = 21999, upload-time = "2026-05-05T18:04:08.971Z" }, +] + +[[package]] +name = "awscrt" +version = "0.32.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/cb/980fe60c4209af71d036276217f8b9f372f958e290c15d2849a3de4dcd23/awscrt-0.32.2.tar.gz", hash = "sha256:a4f48805e8a66237923f03b7b692d213994cff42d1ff08125d1d60c74fcaf872", size = 36862073, upload-time = "2026-04-24T22:59:55.835Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/93/08c8dde6e10c2aa5e50c08730e7913fd68cc861536260115bdf109f20bb8/awscrt-0.32.2-cp310-cp310-macosx_10_15_universal2.whl", hash = "sha256:a9ed98e20ca6fcad8ef32e3c6779aaa3526a5ff3f0aa99d59e0deeff59640375", size = 3405893, upload-time = "2026-04-24T22:58:34.276Z" }, + { url = "https://files.pythonhosted.org/packages/b8/d8/f838299003df9c77b9d582411f009bf2e5bc07fc68566685b434de249755/awscrt-0.32.2-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:acb707df280079d21d20ad1735b38a80d4939133cbaecfc2e4927dd8fc0190bc", size = 3946562, upload-time = "2026-04-24T22:58:37.467Z" }, + { url = "https://files.pythonhosted.org/packages/34/95/a9d9ff694e15550a1fa2f83ac9b3b50cc5f809d66dacc57af4d84cc01c7f/awscrt-0.32.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ab06479bcdcb42b2956f59c0e4138049e5b44c885b7584ea05e39cc4d71b1f99", size = 4236332, upload-time = "2026-04-24T22:58:38.789Z" }, + { url = "https://files.pythonhosted.org/packages/24/bd/92fcf514966e7a139b146282396608493b99e25f745b332075ebe18e129b/awscrt-0.32.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:2792fc2877335673058d0b4e8249ef73dc36b22689fda939506aea6eceb42054", size = 3883087, upload-time = "2026-04-24T22:58:40.292Z" }, + { url = "https://files.pythonhosted.org/packages/5c/7d/cf7ec60d4997ee966ca003a3c7bfc556ee34db10f230f03e5608edf022ca/awscrt-0.32.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:0d2e8a13063a3d9b61eee0e2885d133a75b38de3600f6b62437d8559f7d664e3", size = 4118974, upload-time = "2026-04-24T22:58:41.913Z" }, + { url = "https://files.pythonhosted.org/packages/65/bb/769f8923d7d7986e762898fa98a9d16130681a47da2ff668e6fb07bae0b4/awscrt-0.32.2-cp310-cp310-win32.whl", hash = "sha256:4e975223f3af4faa581c733f0fdd316514b987c42ce85ef3bcd6d9c02eb48f77", size = 4067395, upload-time = "2026-04-24T22:58:43.308Z" }, + { url = "https://files.pythonhosted.org/packages/83/42/17ab1d701f2628adf7b1b993cae9993111ed760aaf08f5772c2bf37c6b7a/awscrt-0.32.2-cp310-cp310-win_amd64.whl", hash = "sha256:a13c0a555bd930c829c72e6b2b2df70442b3b414037fd488984495b5beff5ddc", size = 4224532, upload-time = "2026-04-24T22:58:45.039Z" }, + { url = "https://files.pythonhosted.org/packages/03/99/a6c712a2f638c766b0879da0eacd9fe5695c64deb89c0357bcc4f1f4df9b/awscrt-0.32.2-cp311-abi3-macosx_10_15_universal2.whl", hash = "sha256:32785f54d0786e07b6491b51f9c2f75ea9e17decd39bb6b66fdc60cd871a49ef", size = 3406247, upload-time = "2026-04-24T22:58:46.268Z" }, + { url = "https://files.pythonhosted.org/packages/a7/03/665c81f3b9d3c56fcdfb8f353c22501bc538d192c431cfaaf307f867d404/awscrt-0.32.2-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f8a586c52f41ff14c2f1b8afeb764e231ad3d66acfd42a6b9fb6c8afd8da8fe2", size = 3906890, upload-time = "2026-04-24T22:58:47.96Z" }, + { url = "https://files.pythonhosted.org/packages/31/ad/5db0691a0c72d55a17c03c47149cf15d4602b83b470355c322c9a7f115e8/awscrt-0.32.2-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3ef1664588d35dcad2115377120667e689cd7a517da52a481373c9536811ed96", size = 4196631, upload-time = "2026-04-24T22:58:49.244Z" }, + { url = "https://files.pythonhosted.org/packages/9a/66/63a4654b2996158f33e88d74d79178a5812d94a7e0d6c94ec475115dff18/awscrt-0.32.2-cp311-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:7ed7c209136650fe25659436bb4150e5af6eb43d71a0bf294f0bf414428736ee", size = 3818090, upload-time = "2026-04-24T22:58:50.657Z" }, + { url = "https://files.pythonhosted.org/packages/cd/26/b8fee227918465830a0bbba48f54ebf0a5029c3a7d11d4fa973838a79262/awscrt-0.32.2-cp311-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:cfd437122d5a2ec7eb9fffaf2cd8b96543d4a0d7e906b9515b79672005a1607a", size = 4056453, upload-time = "2026-04-24T22:58:52.373Z" }, + { url = "https://files.pythonhosted.org/packages/e3/26/f5b9e9677bf90d1840d4db1513ba92e73601ef82e61a16e747a83493d35c/awscrt-0.32.2-cp311-abi3-win32.whl", hash = "sha256:5197d1e4e780d755632f79f8d32a09a30a9101cccb51ca1694fe25c711b2e801", size = 4066027, upload-time = "2026-04-24T22:58:53.752Z" }, + { url = "https://files.pythonhosted.org/packages/57/0b/fd1798551ecdc8a28d61ecdeda248f99faa3457f83aefa10ab797f466889/awscrt-0.32.2-cp311-abi3-win_amd64.whl", hash = "sha256:80420aa19c074a4c0335f2bd0e4aee3381fa452328d937795a1e0c779f0c052d", size = 4221984, upload-time = "2026-04-24T22:58:55.115Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a3/075e29b39945f398adfb5d5ee4d533e00d37c0f5c68d79b250a1bdb087ba/awscrt-0.32.2-cp313-abi3-macosx_10_15_universal2.whl", hash = "sha256:d2f7aee3bce261ab1ceba1fac404de4d496aa866237161d4257cef92bff9d828", size = 3404901, upload-time = "2026-04-24T22:58:56.492Z" }, + { url = "https://files.pythonhosted.org/packages/e3/41/1c4783b32bf4ec7383156787570ea1221c95c037c2c0f11cbc9e9529ba48/awscrt-0.32.2-cp313-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ff0fff9c2b613d7fabc298b0fd81f0d7056353f3d20271a852a719c5b2f7ccf4", size = 3897627, upload-time = "2026-04-24T22:58:58.144Z" }, + { url = "https://files.pythonhosted.org/packages/af/2d/5736128847ff4a877d50720ea7a48d7e50a56b78741919e4b6ffabffb1ad/awscrt-0.32.2-cp313-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:79bb11d5d1dfdcfd867aac4a026bee11afbd2154279e12b66588442d8c14bdf7", size = 4189579, upload-time = "2026-04-24T22:58:59.56Z" }, + { url = "https://files.pythonhosted.org/packages/90/f5/064b23d4c927e9b63725ee60c88edb75a1e8f5fd1e97d56d4d6a63fe954b/awscrt-0.32.2-cp313-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:47330f421948207a122092042f235cf82d48fa145c446ff4db12cc8cd3a418b6", size = 3809647, upload-time = "2026-04-24T22:59:00.884Z" }, + { url = "https://files.pythonhosted.org/packages/f5/36/f14ea531cccb6ff85c4d50c9e79e61030675520a393b4af23685d24ea018/awscrt-0.32.2-cp313-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:5ec4d8200eb0158b60e35fbb65faa96ef1eb7763f6ce1192827886ee24c6df99", size = 4051532, upload-time = "2026-04-24T22:59:02.593Z" }, + { url = "https://files.pythonhosted.org/packages/64/aa/b12e721c8148a1bbc58a22cbd204d88d0a6263331049d40c057f9dcd3e11/awscrt-0.32.2-cp313-abi3-win32.whl", hash = "sha256:a81f30a501d2eb6ba52c769cd6ecb3f7005512fba4a533211dc717e1115b0d94", size = 4062484, upload-time = "2026-04-24T22:59:04.251Z" }, + { url = "https://files.pythonhosted.org/packages/5c/26/9f5f23647465a4fe1b28afce334e54a4264e23b69365024c28955f0c4119/awscrt-0.32.2-cp313-abi3-win_amd64.whl", hash = "sha256:8d731edda20dce5afc15a04731d91136a31779a244672d1f0a292a8b04aa0fd3", size = 4220425, upload-time = "2026-04-24T22:59:05.627Z" }, + { url = "https://files.pythonhosted.org/packages/dd/76/18077410c1d7b524a7a7486550bc969218f6575288f66e285157dc78e143/awscrt-0.32.2-cp313-cp313t-macosx_10_15_universal2.whl", hash = "sha256:f2a085d0dd4eef974f2ae5467ae3717d2fc08dd8cd508c4fd7a5acb658c68616", size = 3414715, upload-time = "2026-04-24T22:59:07.01Z" }, + { url = "https://files.pythonhosted.org/packages/d3/12/465d72ea6afffc7829f7f48f408a707d4dab9e3f2b694f4e0e63e011478a/awscrt-0.32.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d8a4e52f7312e5e435461119aa903f6424e9996d93a040101fb1eb7b9c4e58dc", size = 4028634, upload-time = "2026-04-24T22:59:08.303Z" }, + { url = "https://files.pythonhosted.org/packages/e0/ee/2e9d3716cf6a24f30b85af24691d473c913c119dc996bcd8c9b2b3fe2c17/awscrt-0.32.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6cb3c858e96ba023de691a3b6478bed9fb59085433042dff78c42e59eed19cf2", size = 4311846, upload-time = "2026-04-24T22:59:10.02Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ce/aa08aad92e48929cfe243ed7da5cf039fbb137c391e84af79e88c848a37a/awscrt-0.32.2-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:7f92b8f40a104a2a87ea5f428b3799220666ec1450b3a90665867d3715749e91", size = 3952242, upload-time = "2026-04-24T22:59:11.632Z" }, + { url = "https://files.pythonhosted.org/packages/42/e3/09ad98aa7dae9cd3ad578c7721b64e9da03f9a5ed444e518c63e82db05a9/awscrt-0.32.2-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:9fd2dbca02f4e22fb5c02d1505327e6e6e9320dbe8ca80fc033cbbb29ed8631a", size = 4187982, upload-time = "2026-04-24T22:59:14.477Z" }, + { url = "https://files.pythonhosted.org/packages/26/be/897b1ad519d83a837d721f4e8a87d98e6f36e5b985fa697f3ffed512a8eb/awscrt-0.32.2-cp313-cp313t-win32.whl", hash = "sha256:023a2f4595804a0f1d61ab49b64dda5612be9bfbe9b13759331e8e31658dda3f", size = 4111958, upload-time = "2026-04-24T22:59:15.857Z" }, + { url = "https://files.pythonhosted.org/packages/71/54/6cca298e8acb6769c8078b39bedbc507f17a3f7fe6ea768d8256d584d4ed/awscrt-0.32.2-cp313-cp313t-win_amd64.whl", hash = "sha256:a2f513f0fb3aafdf7cdf29d7f6f0c46bf4cc7880380c86e88635b7818565e76d", size = 4271679, upload-time = "2026-04-24T22:59:17.425Z" }, + { url = "https://files.pythonhosted.org/packages/c7/09/51fca333d42b53ddb04881f53e3cc0f4462872ac81426fbb34f5d0b1d1fb/awscrt-0.32.2-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:7404cb551046bbe7cd454ea75f88b4a1b26f018d1b9bea83dbe46c174789d835", size = 3414716, upload-time = "2026-04-24T22:59:18.918Z" }, + { url = "https://files.pythonhosted.org/packages/39/3a/0ae52a8dfb0e3c37f0129232d79307c65e6ee1ea13a3cbdcf301aaad0a6a/awscrt-0.32.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:51d6132e9d70de40da07dfad5f17780a652dd4b351c35ca97c79d0fa0186d645", size = 4028980, upload-time = "2026-04-24T22:59:20.316Z" }, + { url = "https://files.pythonhosted.org/packages/2d/7f/7f52020bfcf29122cad77e936d82abaa1f6857a5a70453e8cc734119f0e2/awscrt-0.32.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:193cc3ecf03a1dd6f989853b6c21549a0a9750c855520fedc8c3c8fbcd32e1bc", size = 4312564, upload-time = "2026-04-24T22:59:21.851Z" }, + { url = "https://files.pythonhosted.org/packages/25/86/2c9b5479e08b8198459d2a781df5524e45d4d0f9c6329bad26979a4d1e85/awscrt-0.32.2-cp314-cp314t-win32.whl", hash = "sha256:25c7e7e6535cb2d2a4d22fd6264f621672d3903491318364dbc59066b63c7186", size = 4195784, upload-time = "2026-04-24T22:59:23.769Z" }, + { url = "https://files.pythonhosted.org/packages/ab/85/a12515514de8969b9e7fa40bb782501a336a8a985cc3093309120a80b627/awscrt-0.32.2-cp314-cp314t-win_amd64.whl", hash = "sha256:a6782c19a00f354c7b232b675f09cde94d1ca37bcf29009b8779b3f6395b27b4", size = 4366435, upload-time = "2026-04-24T22:59:25.53Z" }, +] + [[package]] name = "azure-ai-contentsafety" version = "1.0.0" @@ -3231,6 +3295,92 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" }, ] +[[package]] +name = "ijson" +version = "3.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/06/b31f040a8764336a11152e474a7abcb3782fedb0d1cdf78f442b82878c56/ijson-3.5.1.tar.gz", hash = "sha256:af40bd1a85f55db0b8b30715c858761306bd92d5590148636f75c3309e6e76bd", size = 69913, upload-time = "2026-07-06T17:37:42.923Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/b8/6401c0e2f99aeff22fc740a1b1c2328269a81050c0c178462d0452e27c7e/ijson-3.5.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:8b4ed62287feee41b90b55ae2800ef56d6bdfd2fbfa02b4fd0634cd4524bc995", size = 89054, upload-time = "2026-07-06T17:36:03.274Z" }, + { url = "https://files.pythonhosted.org/packages/a0/ad/8d9e1f076560efcc6727b06f3276f30bb811961332d83567de70c179e0e8/ijson-3.5.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9708c0a3d1f86056049de631933aef8ec57f2008d4cb55ce241790c7ed557428", size = 60674, upload-time = "2026-07-06T17:36:04.326Z" }, + { url = "https://files.pythonhosted.org/packages/ab/e7/8f001e823846c270e0e9c3526ea99dc3b1ba51b9501e060d8337830d6c76/ijson-3.5.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:904e8cf9ca69f5de5b6bb405a4a075ce3da3413ad50c11f6813f1201e14a8e45", size = 60738, upload-time = "2026-07-06T17:36:05.283Z" }, + { url = "https://files.pythonhosted.org/packages/ae/97/c023067cb5ba4cc455a92110a021863fbe3dc3ffcca34ef95aea9290b8f1/ijson-3.5.1-cp310-cp310-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:8cb5db5bc122da64efb24ce358752d5e097ab41d224ce2992536a0f9073fe4fd", size = 126651, upload-time = "2026-07-06T17:36:06.211Z" }, + { url = "https://files.pythonhosted.org/packages/a6/93/7c2207377b40bc1227c8fe1811e080f3b73cd4a9486af9c1166486c3156c/ijson-3.5.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cae04eff4006fc36bf0b030b38e2646a97092d87d933d20cfe7262e26ed32321", size = 133200, upload-time = "2026-07-06T17:36:07.239Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ea/e4d3f64822fb29d54970909e1e2784daa17f75fe3c6c27544fe92e247aad/ijson-3.5.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70542d4542f079c394e525559188d69e3ccfbfd9bab899acd0bf1dbc7323ddd5", size = 130361, upload-time = "2026-07-06T17:36:08.332Z" }, + { url = "https://files.pythonhosted.org/packages/03/77/a61b6b68868a7368a0e4335975c5352e6c354d05eb73dbef19e796b3eaab/ijson-3.5.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:1321495807dcdaca002cb45f24033208ce1d9f5ffc0c5a5584c5f466d0dcbbd5", size = 133618, upload-time = "2026-07-06T17:36:09.41Z" }, + { url = "https://files.pythonhosted.org/packages/a5/0c/05bde03ef651ae2e1033f136c56f7f5565e9f53e7ff91ca83bfd581cbafa/ijson-3.5.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:9fac9284d62c4317d541274e15a6a6ab6f6d22561579f6570967e3a6eaafaebc", size = 128554, upload-time = "2026-07-06T17:36:10.464Z" }, + { url = "https://files.pythonhosted.org/packages/41/42/29bb5561c60e1f9d58d4fbef686e35b9440d9b56f9254c1c70b807c8f649/ijson-3.5.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:1be3a586c8821ecab9ea8b256f39305c8a0cc33222fe393bcc1fb9221470732b", size = 131233, upload-time = "2026-07-06T17:36:11.783Z" }, + { url = "https://files.pythonhosted.org/packages/69/f7/b0176baac5129b79aa366161d5f524ead91b901f16a5020e495c3f83bcc5/ijson-3.5.1-cp310-cp310-win32.whl", hash = "sha256:3ab6378d9c19f01f206f27f762837ad3979330cabd7864e1b17934c03de6056c", size = 52221, upload-time = "2026-07-06T17:36:12.806Z" }, + { url = "https://files.pythonhosted.org/packages/ec/ef/a707b5830722e9f7af347945f9ee0f360d38922366bc1400c6177154eb9c/ijson-3.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:0663f718c6123899c6bfd9c449ec195cd8c67666b7ea2c7b36fa0cc0dcb13e17", size = 54641, upload-time = "2026-07-06T17:36:13.724Z" }, + { url = "https://files.pythonhosted.org/packages/bb/6b/834e7a4ec7e1019b596daf8d74f697aa1d3e38a17a9c31af6081c070557b/ijson-3.5.1-cp310-cp310-win_arm64.whl", hash = "sha256:0a682954b60fcd0c23d504df6fb1ebde051305e41c9b350f39a3b8bfb168def7", size = 53954, upload-time = "2026-07-06T17:36:14.718Z" }, + { url = "https://files.pythonhosted.org/packages/97/d3/16d1595d3ef4743fc55129211bc52f52d59c582d0b7be045d8c04be0ae0c/ijson-3.5.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2aa9d0cf21d4de89fb633e5ec27e9ad02c3f9a4ffa3940d120b23b8aed3acffc", size = 89069, upload-time = "2026-07-06T17:36:15.727Z" }, + { url = "https://files.pythonhosted.org/packages/32/a5/ddba126e2d46cf3b86ad762aeb5e0a02ce0ebc6e4529fe7d06eecb217844/ijson-3.5.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:05eba5268a38809ba1c3dbfa44ea67336e2c353fc11768acc9c6442fe0ccac50", size = 60697, upload-time = "2026-07-06T17:36:16.66Z" }, + { url = "https://files.pythonhosted.org/packages/dc/74/444d8d00a4506a79fc5544614106fa48d5f6f7049511148d8b6cddb8e9d7/ijson-3.5.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:40ddd236c80a667dd6a1f6b625d18ddac68b8719ff795761b7542f2e1f78e4a4", size = 60747, upload-time = "2026-07-06T17:36:17.927Z" }, + { url = "https://files.pythonhosted.org/packages/ee/b1/bc07831e646aebcc91a7bad9c5a0bf7c3f3395f0b10599e021667a3777f1/ijson-3.5.1-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:e6cf9e49902f28af7a2e2f8b35c201195c0f0d5c170a5786e0c0a1b8492a4e37", size = 132095, upload-time = "2026-07-06T17:36:19.022Z" }, + { url = "https://files.pythonhosted.org/packages/1d/1f/b4547461d75db40744616e40c0a06cf2f46a14e60742f6d12510f4612985/ijson-3.5.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6ee1e6d59c800aa819952f6cb5ff08707ecd576b29cc9c3d00e33c2b371a92ce", size = 138790, upload-time = "2026-07-06T17:36:20.22Z" }, + { url = "https://files.pythonhosted.org/packages/a7/30/7ecba8377509eaea2666db5b39a1a99e23f5e3e1e7ee371ec366cbfc4f7c/ijson-3.5.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:affb85eb75fa03a21d1f790bbf26a0e66e5701672062a30dc5c3c6a29c5c0a63", size = 135233, upload-time = "2026-07-06T17:36:21.252Z" }, + { url = "https://files.pythonhosted.org/packages/38/36/0679010904b24398336b3099b09ccb1daa41c534e7cb0931e89d5fcdbee4/ijson-3.5.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3060b141ef758be3742315d44476109460c265b88247e3a4e479949f8b134eac", size = 138832, upload-time = "2026-07-06T17:36:22.323Z" }, + { url = "https://files.pythonhosted.org/packages/b0/90/a40f971e78191e423c7b3a23756f37c3a51c27aadd7769b3fb1816e0044d/ijson-3.5.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:ffba9bce60be21b496afc67a05ab8e3f431f87f0282fd6ce3c62004c951a1428", size = 133313, upload-time = "2026-07-06T17:36:23.405Z" }, + { url = "https://files.pythonhosted.org/packages/7b/d7/b012c347d3ab011c0c4f7988dc6e85b83eaab59df1aec089f5db0e7b29c5/ijson-3.5.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:170cc4c209f57decc9b7ee5fd340f2a1602d54020fa222846482ff1c99e88fdc", size = 135706, upload-time = "2026-07-06T17:36:24.464Z" }, + { url = "https://files.pythonhosted.org/packages/f5/48/3eacb96124e78271f4e648c6ce36f9ce15ce2cef2afb6f8dc6e213e43979/ijson-3.5.1-cp311-cp311-win32.whl", hash = "sha256:6d581a071dae8dbee61f8d962e892787707bad6e641e2f6fb30dd89d3e896939", size = 52221, upload-time = "2026-07-06T17:36:25.517Z" }, + { url = "https://files.pythonhosted.org/packages/1b/1a/19eff8576da0b46fa4a5c8751536ea27ab34c44b2609b2bcded9d7808d42/ijson-3.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:1356bca96d015948b601b013defb2d5631e4330e8f5880e4d7c933d472a90c34", size = 54641, upload-time = "2026-07-06T17:36:26.453Z" }, + { url = "https://files.pythonhosted.org/packages/c7/80/86b28f28ebf190fffd4f46790e065311e2758b55d8e6bbd33d92e9a49448/ijson-3.5.1-cp311-cp311-win_arm64.whl", hash = "sha256:c2b83b24be73f0c7a301807a4c3081939524421c7ae1556eb6eac7cff50ddfa7", size = 53954, upload-time = "2026-07-06T17:36:27.542Z" }, + { url = "https://files.pythonhosted.org/packages/5b/6e/f3ded1ebb85ccc89a30f7b10a0076f30db70ae1d1e0b6423ff93c57b7539/ijson-3.5.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ee60c7741012671867678eae71c51872cac938b76f3d4ca40a778e6c361774d2", size = 88643, upload-time = "2026-07-06T17:36:28.529Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f2/18f14a1d79ef4898e746b4f50dcdbe60abab317cc2bd8390f043b9553c4e/ijson-3.5.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:11c1d7d36a13054b5872ecd5d745dc4009d9abdbcba2312de69e66c2f92a46d2", size = 60611, upload-time = "2026-07-06T17:36:29.597Z" }, + { url = "https://files.pythonhosted.org/packages/30/c7/6e3e591324fd4c7a7a9e1bc23548bacbd84c0d91766b71f09f13e945e7e9/ijson-3.5.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b9517efbe6604bce16f3e50d49b0cd1bdc58917f98cf2eab026599c5c0422991", size = 60447, upload-time = "2026-07-06T17:36:30.747Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a5/9af7be670381ddac26dd55107ed0110b50f5161673b053311db67f510dcc/ijson-3.5.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:ea4fd7bec203a600b1cc88a492dfe6b75ce4b1b87488a66adcd5406022213f64", size = 139092, upload-time = "2026-07-06T17:36:31.749Z" }, + { url = "https://files.pythonhosted.org/packages/41/fb/f9c1664d75467453e6bd4e5f9cd2211b730b09e049445ab64cbac68cc6a3/ijson-3.5.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:350caea815e53151994b597abc80cf669454276b5ac6aadcec69ef6d48f7e90b", size = 149921, upload-time = "2026-07-06T17:36:32.912Z" }, + { url = "https://files.pythonhosted.org/packages/43/80/d20b1c49c4aa7cc6644131e2e57192b45346ef4816566ed1cd9fd05bae38/ijson-3.5.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e4fcebfe1685bb7ba06a8255a5d428ea6b4b895d7acf979cb637d8bbc9db2f47", size = 149848, upload-time = "2026-07-06T17:36:34.032Z" }, + { url = "https://files.pythonhosted.org/packages/fd/fc/5baa710869f5ab939e6233583ced1546889b55c35f35b844c518ac10abc3/ijson-3.5.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d78f362f51c8691798758a9e6ac3c9d385ee1228cb82987c91562a2fae235cd3", size = 150810, upload-time = "2026-07-06T17:36:35.19Z" }, + { url = "https://files.pythonhosted.org/packages/54/16/a12b3d987a5c1677b04557c6f9b9feb7e04b7d4171e9a344856cb9136e9b/ijson-3.5.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:0b184180d45f85fd4479659582749b109e49f4a29c21ac700ccc9c2280fe015e", size = 142989, upload-time = "2026-07-06T17:36:36.23Z" }, + { url = "https://files.pythonhosted.org/packages/ed/63/1026c535671fc334fc85aeb78f0945c825e7a338575edc753c0f455459ae/ijson-3.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e353891d33a2e6aa5caf72c2a5fbadd7a46f5f9b32dcfd0c84113b2444c255b8", size = 151702, upload-time = "2026-07-06T17:36:37.296Z" }, + { url = "https://files.pythonhosted.org/packages/cb/af/b58aa3a2bf4d31c388ea78b49826605f60932891ce97e404d196766b4ea3/ijson-3.5.1-cp312-cp312-win32.whl", hash = "sha256:936f28671f018f8ac4d3f003ae9fa01d0467ab4ef4cfd0c97f23beda485b61c6", size = 52613, upload-time = "2026-07-06T17:36:38.345Z" }, + { url = "https://files.pythonhosted.org/packages/04/66/ce70a92949c2a753dad91fdd5761dc14f3a44517e80cfc3c26612982ed61/ijson-3.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:322c783f3ee0c6b383bbd4db88370b10172168808cc2a0bf811f1253f7435602", size = 54729, upload-time = "2026-07-06T17:36:39.337Z" }, + { url = "https://files.pythonhosted.org/packages/a5/ff/e17784240c9cf1d58de2f2853ebaf9cc54f6bce117a1f12a6150bbb4a5aa/ijson-3.5.1-cp312-cp312-win_arm64.whl", hash = "sha256:e2ac204b59f09e38e16d277f906240e9fd38780e42076599419265af183dc4b4", size = 53714, upload-time = "2026-07-06T17:36:40.308Z" }, + { url = "https://files.pythonhosted.org/packages/fd/c0/5384ccf4fc497ae3dc79a5a28561b05518b503ade29daf3898168d640406/ijson-3.5.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:3c0556d628443d3e871f414855313b2ae6cd9faa0104de3316bd8db03aab1589", size = 88652, upload-time = "2026-07-06T17:36:41.278Z" }, + { url = "https://files.pythonhosted.org/packages/8e/42/58769b8b6d614adb15c2c938c77bcdbfadfba8b1d21a98b5b09cb8961adc/ijson-3.5.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:12aa7fcf46f0fdc8e9e7cf37541e1dc20ac3f9243a23f4d346ab5395f72b0fe2", size = 60607, upload-time = "2026-07-06T17:36:42.697Z" }, + { url = "https://files.pythonhosted.org/packages/db/4a/8322c2824c24184880587bbca45531127a21a4b3bfc897f13427fea02424/ijson-3.5.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a96066d8c12a18ce2fa90579f2bbf991377cb71725874932e4a5d855226c162a", size = 60447, upload-time = "2026-07-06T17:36:43.791Z" }, + { url = "https://files.pythonhosted.org/packages/f4/43/7bdca8f733c45ce97f61a64fadd3e51d255c4c9b467345cbf71ccc7bb368/ijson-3.5.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:a19413a092d458a57aaa574fec08e265851d3b5c6e018377f426cd5e70b91280", size = 138889, upload-time = "2026-07-06T17:36:45.081Z" }, + { url = "https://files.pythonhosted.org/packages/e7/dc/e8a2e63700ab1d63aaf3fa38c454f8178eaa5b80a6d7c019d1d61b490a6c/ijson-3.5.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:65974568748678165d7e90e3e7ce2f7c233cfe4de6c37fbb0760941c97e14632", size = 149933, upload-time = "2026-07-06T17:36:46.312Z" }, + { url = "https://files.pythonhosted.org/packages/d9/56/640a4d980f7f2c11e399a7fd5ccb9e3d3c9e1dec3a1d5a10024570697c25/ijson-3.5.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bad5d55c99c89de8cd0a4cded51f86427ba3353c4dccca37ec2e32e06f26b437", size = 149857, upload-time = "2026-07-06T17:36:47.309Z" }, + { url = "https://files.pythonhosted.org/packages/3d/a1/c953e22c83992b69ae538a83b3678d28768f1a48042fc7794733423a5ce7/ijson-3.5.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1a38d503ce343952e88edfd9a27296a4ec96af7073a9db58b3df6233367f75fc", size = 151141, upload-time = "2026-07-06T17:36:48.405Z" }, + { url = "https://files.pythonhosted.org/packages/9e/ab/8fe5b7269b140e6e5f8837a33ce980fd9b67c70d0f8114289ed1cea4dace/ijson-3.5.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:2f41982c73896acab4a2a14faa14e152e444bd69f37c3139204429fd3fe65a10", size = 143112, upload-time = "2026-07-06T17:36:50.353Z" }, + { url = "https://files.pythonhosted.org/packages/78/f3/23d1284edcde50ba337ddfba5b5d59f8273084d98b28af94715e73dd2b64/ijson-3.5.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:3321fede2b638d400de0036889a3a25c3bb689feb8df45e70a393346aad6194f", size = 152184, upload-time = "2026-07-06T17:36:51.536Z" }, + { url = "https://files.pythonhosted.org/packages/82/4e/df61be89dd295e4da722ec96ba03b1765bcb2becdaaaede9c96a7d2365b6/ijson-3.5.1-cp313-cp313-win32.whl", hash = "sha256:af6ddbd10ac9bce87a835f2de3ec61455ec435c54e7e0ba7b17c31c66de6f164", size = 52607, upload-time = "2026-07-06T17:36:52.596Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d9/03e5dbd3ef7e0cee06fbef0f87b91d7ce1c07fae9b5a1b0ca8b895de62c4/ijson-3.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:1de3de278b0ffb40338374ad2a730e1c56f933e0706b1815ebeb07b82239b1a3", size = 54730, upload-time = "2026-07-06T17:36:53.526Z" }, + { url = "https://files.pythonhosted.org/packages/38/30/4f37076c88a96a1a5e44df38b59fade4f59eaef87ef8b5162d55b2d426d5/ijson-3.5.1-cp313-cp313-win_arm64.whl", hash = "sha256:c8a36a19b92cb7172c6448ab94f446033cfa3129dc4894aebe205f96b3fabf42", size = 53719, upload-time = "2026-07-06T17:36:54.592Z" }, + { url = "https://files.pythonhosted.org/packages/f9/17/54f9180c0da9a9e96e5b3791bc74093f029a2344678b4da218c2699465bf/ijson-3.5.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:21e1a250b254edba2f0dd7272a4c56f0a879aabe328d9e306dd1fc115f560e74", size = 89223, upload-time = "2026-07-06T17:36:55.534Z" }, + { url = "https://files.pythonhosted.org/packages/09/70/0ee0d2627c534174455a745ca25284797e71b0d6e2b2a1b31cc914e7b462/ijson-3.5.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:e01f95433725e2df62d682ff88e4a57bb694385ff2362bc364adec961167ae04", size = 60831, upload-time = "2026-07-06T17:36:56.554Z" }, + { url = "https://files.pythonhosted.org/packages/8d/e6/56f64ba7a3e7a25d9a9fbbeb4c30597d6b76c1094cc2041d11a3224b562c/ijson-3.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:539e8d6cca079bcbb68c390e55148f908e0a943a34f7dd321248637c6272adca", size = 60752, upload-time = "2026-07-06T17:36:57.826Z" }, + { url = "https://files.pythonhosted.org/packages/3e/2b/5a55db881f1b043cd6d5716578937a60ac16348be1a3afbf846b21cf4b44/ijson-3.5.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:32f64051be2f990d8ae7b614b5abdf4a7bead510ce3666568d7403c6c46ce4d8", size = 140783, upload-time = "2026-07-06T17:36:58.984Z" }, + { url = "https://files.pythonhosted.org/packages/2e/61/f7783cc18672dc31544141139efd187fb34795d24e573fed6abea6b776c7/ijson-3.5.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cd0dfc5a788d0b0c2f1eab258b9dabdeefc631ca8ef87644a999f633b0b2555a", size = 149976, upload-time = "2026-07-06T17:37:00.235Z" }, + { url = "https://files.pythonhosted.org/packages/5f/d6/4182dd63b6b70eae4f5208c53558a050895a40734dff283463033c153742/ijson-3.5.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42bfda7858d99ee9777ec28cb6d347928249eefeb577f9b0a67503c18f7ebb6a", size = 149317, upload-time = "2026-07-06T17:37:01.476Z" }, + { url = "https://files.pythonhosted.org/packages/01/b1/a675e4a9b428a0ef556e7d718bf0e6885e3e5543042248a1a7030899a3d4/ijson-3.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c4b9a28e9719d1aebebe93ad8dc2ba87f4e2d9035043b196c1c07ef8530b44cc", size = 150555, upload-time = "2026-07-06T17:37:02.676Z" }, + { url = "https://files.pythonhosted.org/packages/b5/69/52686f56b44af63a93c3dc3f5bcfa07f87427d9aea4d2cbe3e1c94188c74/ijson-3.5.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:9a0b25c750a6bde14a0b31f1dcbfc86368e50767e3eaa73bb138e54128055edd", size = 144485, upload-time = "2026-07-06T17:37:03.779Z" }, + { url = "https://files.pythonhosted.org/packages/f0/46/10554e817dde56300a8414e52c0f5a44a29f3440327cd6d829ece57759b3/ijson-3.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:bd756f7b22df745ac14b7bc2ab9ed7c190a222e4c8e1bef26ef1162af8e54d0f", size = 151470, upload-time = "2026-07-06T17:37:04.901Z" }, + { url = "https://files.pythonhosted.org/packages/91/82/f37cbb110b48abdb623d169d0e196f2f6e064e2c20fa789ecde6e69b0440/ijson-3.5.1-cp314-cp314-win32.whl", hash = "sha256:e035cdfb2a1446b13881f0dfc0eecd1541cbb17a27a938ded2160ae6ce25051b", size = 53219, upload-time = "2026-07-06T17:37:06.254Z" }, + { url = "https://files.pythonhosted.org/packages/00/58/792df8f001c246c8ff28f860de81d35ea0d797c0d3276c22a2af83089656/ijson-3.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:eeb2fb2daa5dd30326f93db465d0855b34aa6b1f52a7c0ff94522aec5ad57dfb", size = 55485, upload-time = "2026-07-06T17:37:07.242Z" }, + { url = "https://files.pythonhosted.org/packages/c0/3c/db3ccc22c09ed4738787e8d82fff76101aa81ec8de7eaf6572e065e012d3/ijson-3.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:a96ab35d7ce2129dfde49c4c807596443410e260d7f7a4ca8fe4d0035553b589", size = 54390, upload-time = "2026-07-06T17:37:08.497Z" }, + { url = "https://files.pythonhosted.org/packages/26/59/eefa5d9488250c03f24152576804205ae40e29cac0dc65cbbc5f3d422008/ijson-3.5.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:77b68e91f95fb16ac2e7819903cd545db6cffa308c28833cc34911e6b21e91dd", size = 93177, upload-time = "2026-07-06T17:37:09.71Z" }, + { url = "https://files.pythonhosted.org/packages/88/db/6329eb7bb9f1906c1906fc10e7074b8f08bf39b7d50baa58f1b597d48898/ijson-3.5.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:94a95065b1ac67602af0cec852b07505abc37b77e3774d1c801d935d05e48f82", size = 62891, upload-time = "2026-07-06T17:37:10.735Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d0/b3beddb96eef0b20bb9902c36e4de30f145be06d7e5e1d780e1a1689d0ce/ijson-3.5.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b70b5da6b0571da8f601a437c4fba2d35bc27739637d85f3acdc8f88916ce68e", size = 62575, upload-time = "2026-07-06T17:37:11.681Z" }, + { url = "https://files.pythonhosted.org/packages/5b/01/95f3a7c27d25bb917954ef0c8e86d0e60f585b9db675cbd05d355f54cce8/ijson-3.5.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:0ade373dd765b057b1dec05d7711bfeb5a36f1e825259466d9f545cfd8ef3ba3", size = 200568, upload-time = "2026-07-06T17:37:12.743Z" }, + { url = "https://files.pythonhosted.org/packages/77/61/c94ee4ea1f22318aab9a49b35d0ce8ac87dd24d508ea4c77dcbde362ba5e/ijson-3.5.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:882bc0bdd25d41eae90a15695cd50707edde0978b8b72a2532e30442dd8fd04c", size = 217956, upload-time = "2026-07-06T17:37:14.041Z" }, + { url = "https://files.pythonhosted.org/packages/1a/82/43e8d225aea5ee00eef7998c8ce41f344f7ba451329dfa9e92f4700813af/ijson-3.5.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:451901c36e12fa87cbb1cafe661bd25c08c6bd7900cc738279614f71cea07048", size = 208403, upload-time = "2026-07-06T17:37:15.201Z" }, + { url = "https://files.pythonhosted.org/packages/cf/6f/375f67fad76677aca9bc0817b2b18fdd231d309fe24e26b19a5556ef6cdd/ijson-3.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e3c5f660658f2ebfba5d4dfe4bafe8cd3a0defcda410ec08d2205fe08c398940", size = 211967, upload-time = "2026-07-06T17:37:16.484Z" }, + { url = "https://files.pythonhosted.org/packages/dc/53/4c754c3ba18ec70b7086b91a4abd368358fc47cc9b3871afd50deef4fea1/ijson-3.5.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:29eb8f0c77a296a10843a1714ad4a5d561e604cda3c88585e9012cf2c1729b0a", size = 201020, upload-time = "2026-07-06T17:37:18.017Z" }, + { url = "https://files.pythonhosted.org/packages/26/2d/3e7191b3222a31c378b827565b4fa64676a293441279f84db3d971720bf5/ijson-3.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:85997568d6b304cfa59d5c3f2b04f95b92e9a8c7f57d312343a7989cf8dfff85", size = 205584, upload-time = "2026-07-06T17:37:19.343Z" }, + { url = "https://files.pythonhosted.org/packages/24/11/55ae9c915e68f37c8698f8b09355071dc808ced5e9d4abf8238dc363f500/ijson-3.5.1-cp314-cp314t-win32.whl", hash = "sha256:c2e2509dc7f2fa5a2ac9ba7d15dd901f4093bd36b0784f65e04b681b7956651c", size = 54438, upload-time = "2026-07-06T17:37:20.656Z" }, + { url = "https://files.pythonhosted.org/packages/96/df/5bf2656447f14a923d25a0401b1cd628ca05c23041d3a4c116ae8d44dc39/ijson-3.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:2699e838099d056818c5f8e4ba702b345d0304e58847bdc79c5c1616d5d750a5", size = 56467, upload-time = "2026-07-06T17:37:21.615Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e4/dec06e84fac704039625039c6b116a44f17ad72fda48b8f88a2493364b77/ijson-3.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:c388f85cbb9eec022b2bdedd23ffacfe7ab100c1200b1f47bee6e6ea2c3309fa", size = 55774, upload-time = "2026-07-06T17:37:22.958Z" }, + { url = "https://files.pythonhosted.org/packages/49/ea/f42470cc773c8686dd0823da8aefc31a138cd9aea1ad476d43c8293068da/ijson-3.5.1-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:077b1b0bcb6a622d460c6674fe6647c7af5a3b06503e1996d1efcf9f78c94512", size = 57830, upload-time = "2026-07-06T17:37:37.005Z" }, + { url = "https://files.pythonhosted.org/packages/d0/2f/64c61edab2c5ecf42a524146a70fa6171c8cf3960b947fb4c5f175660cb3/ijson-3.5.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:e8dbf71b21e65cb7f0d4d387c07fe73be820168070c3be05a0763a80f424f1c7", size = 57325, upload-time = "2026-07-06T17:37:38.017Z" }, + { url = "https://files.pythonhosted.org/packages/9f/5b/553ea8f14dfc756d6b6c9be2e2231ab44877ce96408eb9da3bb3f11ddd13/ijson-3.5.1-pp311-pypy311_pp73-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:0d7c5025a820f36f3e0e64f4b0232b338c690664c12b497e205cf64dcc64fc12", size = 71344, upload-time = "2026-07-06T17:37:38.997Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3e/0248fd00746731074ca01365a25d8aa3c4d54642c8a14490d94f7550bda9/ijson-3.5.1-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa7a2c94e43c02e0482088e6ff997e2bd7b9a76e6f1d0fd70891b4b5ff51318f", size = 71335, upload-time = "2026-07-06T17:37:39.965Z" }, + { url = "https://files.pythonhosted.org/packages/7a/b9/1f1259546cc875adad240c468515f428d3a79b3def3ced17be3cdfe29146/ijson-3.5.1-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:69b5eef70240e9734c5a2fb5cc3742cae411fc833a66b9a50722b9eedb1e27de", size = 68728, upload-time = "2026-07-06T17:37:40.928Z" }, + { url = "https://files.pythonhosted.org/packages/ea/02/aafbf0c3e1468c7c0f607065363b49c381de7e4bb43ae6674684a3fafe92/ijson-3.5.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:4b75b6bf4b0dbb0df24947db6722cd5723ce8d6e6b13fddbfc98db312ba82237", size = 54922, upload-time = "2026-07-06T17:37:41.879Z" }, +] + [[package]] name = "imagesize" version = "2.0.0" @@ -3984,6 +4134,9 @@ dependencies = [ ] [package.optional-dependencies] +bedrock-realtime = [ + { name = "aws-sdk-bedrock-runtime", marker = "python_full_version >= '3.12'" }, +] caching = [ { name = "diskcache" }, ] @@ -4175,6 +4328,7 @@ requires-dist = [ { name = "apscheduler", marker = "extra == 'proxy'", specifier = ">=3.11.2,<4.0" }, { name = "audioread", marker = "extra == 'stt-nvidia-riva'", specifier = ">=3.0.1" }, { name = "aurelio-sdk", marker = "python_full_version < '3.14' and extra == 'semantic-router'", specifier = ">=0.0.19,<1.0" }, + { name = "aws-sdk-bedrock-runtime", marker = "python_full_version >= '3.12' and extra == 'bedrock-realtime'", specifier = ">=0.7.0,<0.8.0" }, { name = "azure-ai-contentsafety", marker = "extra == 'proxy-runtime'", specifier = ">=1.0.0,<2.0" }, { name = "azure-identity", marker = "extra == 'extra-proxy'", specifier = ">=1.25.2,<2.0" }, { name = "azure-identity", marker = "extra == 'proxy'", specifier = ">=1.25.2,<2.0" }, @@ -4254,7 +4408,7 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.21.0,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] -provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "proxy-runtime"] +provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] ci = [ @@ -8374,6 +8528,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/0e/3ae19fa941522cd98e119762e7181d371c8dba0b2d72bfaf9522692e329c/skops-0.14.0-py3-none-any.whl", hash = "sha256:60a5db78a9db46ccee2139a0ba13ab5afb1c96f4749b382e75a371291bbe3e36", size = 132198, upload-time = "2026-04-20T18:23:54.018Z" }, ] +[[package]] +name = "smithy-aws-core" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aws-sdk-signers", marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, + { name = "smithy-http", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fc/a8/37bfde59519f45d2047d0033b791aca6574d867aaf57bb56a6de42ab5c26/smithy_aws_core-0.7.0.tar.gz", hash = "sha256:34e82d09fc808acd5ffc80f03828d0609c6a211f49f0884dc6ee7ca095a1b6af", size = 15670, upload-time = "2026-06-23T04:04:50.365Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/54/2d06dd9a3972a380d71bb8c3312e317aa8f1ea68dd28cffc06955ccf0220/smithy_aws_core-0.7.0-py3-none-any.whl", hash = "sha256:6c60c8fbb9431c60e80ea7f2d37e7ae48409cc1541f587fe073f202eca067e92", size = 24894, upload-time = "2026-06-23T04:04:49.349Z" }, +] + +[package.optional-dependencies] +eventstream = [ + { name = "smithy-aws-event-stream", marker = "python_full_version >= '3.12'" }, +] +json = [ + { name = "smithy-json", marker = "python_full_version >= '3.12'" }, +] + +[[package]] +name = "smithy-aws-event-stream" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/0e/6efb3a4ed92c0f1ada6de060ac92e7115a1e34d0ab1fb99a6056734a88ea/smithy_aws_event_stream-0.3.0.tar.gz", hash = "sha256:a0e227367a973144e205a075d0a424f95c92f26656a1018d08900da2ae547c49", size = 12818, upload-time = "2026-05-05T18:04:14.317Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/b4/c4c4a9aa5a4cf3ee285f1bd71999441177c651d911a82a6e64c57b6e3e65/smithy_aws_event_stream-0.3.0-py3-none-any.whl", hash = "sha256:8b505cc28230e4fe9c5e025333209b44ca2db560451ac9ed9a7d74939edf4413", size = 15845, upload-time = "2026-05-05T18:04:13.351Z" }, +] + +[[package]] +name = "smithy-core" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e9/45/688d52c61cd4d843bb230694259e91d4c7d6954eeecbadf452a168001d45/smithy_core-0.6.0.tar.gz", hash = "sha256:ba2e5d860d716aff75004a23f53e09dfaca3e2b94f8a00c1f76dcb355b769ce0", size = 52095, upload-time = "2026-06-23T04:04:44.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/b6/06795faa9844b9667ae492e6293370393e19e7f0c2df8da1b4bf7e5f6ed9/smithy_core-0.6.0-py3-none-any.whl", hash = "sha256:51e347ed309d60ab9d36b783dbf88de614c460d51bec79d39cd403956b00f063", size = 66879, upload-time = "2026-06-23T04:04:43.596Z" }, +] + +[[package]] +name = "smithy-http" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/58/5a772d212e066d6fc1398946c4aae19bcdaa75209879d776f641b6a06b5b/smithy_http-0.4.2.tar.gz", hash = "sha256:50d11b6a55e42448450a01e3d0f605ccee65a72abf52d02eed82862a15be5937", size = 29616, upload-time = "2026-06-23T04:04:45.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/3e/7b2464d40893bec0b5d1f479d25116d4aa09f9f66536b4c4b3126202215d/smithy_http-0.4.2-py3-none-any.whl", hash = "sha256:a158f107e9fab925289d20772c2e38b0bba94e55c05d0edc9290310f22a60454", size = 41025, upload-time = "2026-06-23T04:04:46.764Z" }, +] + +[package.optional-dependencies] +awscrt = [ + { name = "awscrt", marker = "python_full_version >= '3.12'" }, +] + +[[package]] +name = "smithy-json" +version = "0.2.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ijson", marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b9/6c/418b5687d8933b7a135d5e1a98c61fe814b98f72517dbae0e666860cb876/smithy_json-0.2.3.tar.gz", hash = "sha256:686e9b55a36dacb08e472732b358573ef78009055e05e9fce2e806d61490b2b3", size = 7805, upload-time = "2026-06-23T04:04:47.71Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/14/eabb26b355415bcd9feef27fb5b18f1dad3fabd4208cfcbaf152025fa9ae/smithy_json-0.2.3-py3-none-any.whl", hash = "sha256:594e1bbe3d480963237f8fd0fc648dbd4e988b4503fea90157b5f07706796327", size = 10252, upload-time = "2026-06-23T04:04:48.46Z" }, +] + [[package]] name = "smmap" version = "5.0.3" From fa6b20916572821e58618b0e927b2764fceed768 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:56:59 -0700 Subject: [PATCH 034/130] feat(guardrails): add only_scan_new_messages for per-session incremental scanning (#33278) * feat(guardrails): add only_scan_new_messages for per-session incremental scanning Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): use fixed TTL constant and revert unrelated test formatting Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): run only_scan_new_messages in the unified apply_guardrail path The initial wiring lived in BedrockGuardrail.async_pre_call_hook, but the proxy routes Bedrock through the unified apply_guardrail interface, so the flag had no effect live. Move incremental selection into apply_guardrail: filter the flat texts list against per-session scanned hashes, skip the Bedrock call when nothing is new, and mark hashes only after a successful (non-blocked) scan. Full-context fallback is preserved when there is no session id, the cache is unavailable, or a masking guardrail is configured. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover session-id fallbacks and mark_texts_scanned guards Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): fall back to full scan when incremental guardrail masks content, use shared cache Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover generic agent multi-turn incremental scan Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover incremental scan cache resolver fallbacks Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover flag interactions and /v1/messages incremental scan semantics * feat(guardrails): make GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS env configurable * test(guardrails): prove skip_system/skip_tool are enforced upstream of incremental scan --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> Co-authored-by: Yucheng Zhu --- litellm/constants.py | 3 + litellm/integrations/custom_guardrail.py | 102 +++++- .../guardrail_hooks/bedrock_guardrails.py | 93 +++++ .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 12 + .../integrations/test_custom_guardrail.py | 198 ++++++++++ .../test_anthropic_guardrail_handler.py | 129 +++++++ .../test_openai_guardrail_handler.py | 92 +++++ .../test_bedrock_guardrails.py | 340 ++++++++++++++++++ 9 files changed, 969 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2af84c139a1..84ac9e29729 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -264,6 +264,9 @@ MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) +GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS = int( + os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) +) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 856556f7c56..cf9dafcb222 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import hashlib import os import secrets from datetime import datetime @@ -46,7 +47,10 @@ if TYPE_CHECKING: dc = DualCache() -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import ( + GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + PRE_CALL_EXECUTED_GUARDRAILS_KEY, +) from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -113,6 +117,7 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + only_scan_new_messages: bool = False, **kwargs, ): """ @@ -145,6 +150,7 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: ## validate event_hook is in supported_event_hooks @@ -269,6 +275,100 @@ class CustomGuardrail(CustomLogger): """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) + @staticmethod + def _scanned_text_hash(text: str) -> str: + """Stable content hash for a single scannable text segment. + + Hashing the exact text the provider would receive means an edited earlier + segment produces a different hash and gets re-scanned, while an unchanged + segment repeated on a later turn is skipped. + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + def _scanned_texts_cache_key(self, session_id: str) -> str: + return f"guardrail_scanned_texts:{self.guardrail_name}:{session_id}" + + async def filter_new_texts_for_session( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> list[str] | None: + """Return only the text segments not already scanned earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or the cache read failed). ``None`` signals + the caller to fall back to a full scan; a returned list (possibly empty) + signals the caller to scan only that subset and skip masking write-back. + """ + if not self.only_scan_new_messages or not texts: + return None + + if self.mask_request_content or self.mask_response_content: + verbose_logger.warning( + "Guardrail %s: only_scan_new_messages is not supported with masking; scanning full context.", + self.guardrail_name, + ) + return None + + session_id = get_session_id_from_request_data(request_data) + if not session_id: + verbose_logger.debug( + "Guardrail %s: only_scan_new_messages enabled but request has no session id; scanning full context.", + self.guardrail_name, + ) + return None + + try: + cached: object = await cache.async_get_cache(key=self._scanned_texts_cache_key(session_id)) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must fall back to a full scan + verbose_logger.warning( + "Guardrail %s: failed to read scanned-message cache (%s); scanning full context.", + self.guardrail_name, + e, + ) + return None + + seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set() + return [text for text in texts if self._scanned_text_hash(text) not in seen] + + async def mark_texts_scanned( + self, + texts: list[str] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> None: + """Record the hashes of all text segments present on a successful (non-blocked) scan. + + Called only after the guardrail allows the request, so a blocked segment is + never marked scanned and will be re-checked if the client retries. + """ + if not self.only_scan_new_messages or not texts: + return + if self.mask_request_content or self.mask_response_content: + return + session_id = get_session_id_from_request_data(request_data) + if not session_id: + return + + cache_key = self._scanned_texts_cache_key(session_id) + current_hashes = [self._scanned_text_hash(text) for text in texts] + try: + existing: object = await cache.async_get_cache(key=cache_key) + existing_hashes: list[str] = [str(h) for h in existing] if isinstance(existing, list) else [] + merged: list[str] = list(dict.fromkeys(existing_hashes + current_hashes)) + await cache.async_set_cache( + key=cache_key, + value=merged, + ttl=GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + ) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must not block the request + verbose_logger.warning( + "Guardrail %s: failed to persist scanned-message cache (%s); next call will re-scan.", + self.guardrail_name, + e, + ) + def should_route_on_sensitive_data(self) -> bool: """ Returns True if this guardrail is configured to route requests diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 54156715da8..cec682d772a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2046,6 +2046,90 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 verbose_proxy_logger.debug("Applied masking to choice text content") + @staticmethod + def _incremental_scan_cache() -> DualCache: + """Resolve the cache used to remember which segments a session already scanned. + + Prefers the proxy's shared cache (``internal_usage_cache.dual_cache``), which is + backed by Redis when the deployment configures it, so incremental state is shared + across proxy instances. Falls back to a process-local ``DualCache`` singleton when + the proxy is not running (e.g. unit tests), where sharing does not apply. + """ + from litellm.integrations.custom_guardrail import dc as fallback_cache + + try: + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging + except Exception: # noqa: BLE001 # proxy not importable outside the server; use local fallback + return fallback_cache + if _proxy_logging is not None: + return _proxy_logging.internal_usage_cache.dual_cache + return fallback_cache + + def _bedrock_response_has_masked_output(self, response: BedrockGuardrailResponse) -> bool: + """Return True if the guardrail rewrote (masked/anonymized) any scanned text. + + Bedrock returns non-empty ``output``/``outputs`` text only when it changed the + content; an ``action == "NONE"`` response leaves both empty. + """ + for field in ("output", "outputs"): + items = response.get(field) or [] + if any(isinstance(item, dict) and item.get("text") for item in items): + return True + return False + + async def _apply_incremental_request_scan( + self, + texts: list[str], + inputs: "GenericGuardrailAPIInputs", + request_data: dict, + ) -> Optional["GenericGuardrailAPIInputs"]: + """Scan only the text segments not already seen earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or cache unavailable) or when the guardrail + turns out to mask content, telling the caller to run the normal full scan. + Otherwise scans only the new segments and skips the Bedrock call entirely + when nothing is new. Incremental mode is for blocking/detection guardrails + only: if the guardrail returns masked output it cannot be applied to the + skipped context, so the scan falls back to the full path and no session + state is recorded. + """ + cache = self._incremental_scan_cache() + + new_texts = await self.filter_new_texts_for_session( + texts=texts, + request_data=request_data, + cache=cache, + ) + if new_texts is None: + return None + + if not new_texts: + verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call") + return inputs + + bedrock_response = await self.make_bedrock_api_request( + source="INPUT", + messages=[ChatCompletionUserMessage(role="user", content=text) for text in new_texts], + request_data=request_data, + logging_event_type=GuardrailEventHooks.pre_call, + ) + + if self._bedrock_response_has_masked_output(bedrock_response): + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: guardrail returned masked/anonymized content; " + "only_scan_new_messages cannot apply masking to skipped context, falling back to a full-context scan", + self.guardrail_name, + ) + return None + + await self.mark_texts_scanned( + texts=texts, + request_data=request_data, + cache=cache, + ) + return inputs + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -2077,6 +2161,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: verbose_proxy_logger.debug(f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)") + if input_type == "request": + incremental_result = await self._apply_incremental_request_scan( + texts=texts, + inputs=inputs, + request_data=request_data, + ) + if incremental_result is not None: + return incremental_result + masked_texts = [] selection = self._select_messages_for_apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 14e76a21093..e909c15382b 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -35,6 +35,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_sts_endpoint=litellm_params.aws_sts_endpoint, aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, + only_scan_new_messages=litellm_params.only_scan_new_messages or False, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 47d93fc2d7a..c86794b90f8 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -725,6 +725,18 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)", ) + only_scan_new_messages: Optional[bool] = Field( + default=False, + description=( + "When True, the guardrail only scans messages that have not already been scanned " + "earlier in the same session (identified by litellm_session_id / session_id). " + "Message content is hashed per session and cached; only the diff (new or edited " + "messages) is sent to the guardrail provider on follow-up calls. Falls back to a " + "full scan when the request has no session id or the cache is unavailable. Intended " + "for blocking/detection guardrails; not applied when mask_request_content is set." + ), + ) + skip_system_message_in_guardrail: Optional[bool] = Field( default=None, description=( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 9289dece83f..64813c1eda7 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1716,3 +1716,201 @@ class TestApplyGuardrailStyleDeploymentDispatch: await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert guardrail.apply_called is False + + +class TestOnlyScanNewMessages: + """Incremental guardrail scanning: only send text segments not already scanned this session.""" + + def _guardrail(self, **overrides): + params = dict(guardrail_name="test-guard", only_scan_new_messages=True) + params.update(overrides) + return CustomGuardrail(**params) + + def _cache(self): + from litellm.caching import DualCache + + return DualCache() + + @pytest.mark.asyncio + async def test_disabled_returns_none(self): + guardrail = self._guardrail(only_scan_new_messages=False) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_session_id_fails_safe_to_full_scan(self): + guardrail = self._guardrail() + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"metadata": {}}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_masking_guardrail_not_supported(self): + guardrail = self._guardrail(mask_request_content=True) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_cache_read_failure_fails_safe_to_full_scan(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_get_cache = AsyncMock(side_effect=RuntimeError("redis down")) + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], + request_data={"litellm_session_id": "s1"}, + cache=cache, + ) + assert result is None + + @pytest.mark.asyncio + async def test_dedupes_previously_scanned_texts(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-dedupe"} + turn1 = ["you are helpful", "first question"] + + first = await guardrail.filter_new_texts_for_session(texts=turn1, request_data=request, cache=cache) + assert first == turn1 + await guardrail.mark_texts_scanned(texts=turn1, request_data=request, cache=cache) + + turn2 = turn1 + ["an answer", "second question"] + second = await guardrail.filter_new_texts_for_session(texts=turn2, request_data=request, cache=cache) + assert second == ["an answer", "second question"] + + @pytest.mark.asyncio + async def test_no_new_texts_returns_empty(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-empty"} + texts = ["only message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_modified_earlier_text_is_rescanned(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-edit"} + original = ["original"] + + await guardrail.filter_new_texts_for_session(texts=original, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=original, request_data=request, cache=cache) + + edited = ["original EDITED"] + result = await guardrail.filter_new_texts_for_session(texts=edited, request_data=request, cache=cache) + assert result == edited + + @pytest.mark.asyncio + async def test_blocked_scan_does_not_persist_hashes(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-blocked"} + texts = ["please block me"] + + filtered = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert filtered == texts + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == texts + + @pytest.mark.asyncio + async def test_scanned_hashes_written_with_fixed_ttl(self): + from unittest.mock import AsyncMock + + from litellm.constants import GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + request = {"litellm_session_id": "sess-ttl"} + + await guardrail.mark_texts_scanned(texts=["a", "b"], request_data=request, cache=cache) + + cache.async_set_cache.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["ttl"] == GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_session_id_from_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"metadata": {"session_id": "sess-meta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_session_id_from_litellm_metadata_is_used_for_dedupe(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_metadata": {"session_id": "sess-lmeta"}} + texts = ["shared message"] + + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) + + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_mark_texts_scanned_disabled_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(only_scan_new_messages=False) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_masking_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail(mask_request_content=True) + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_without_session_does_not_persist(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"metadata": {}}, cache=cache) + cache.async_set_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_mark_texts_scanned_survives_cache_write_failure(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock(side_effect=RuntimeError("redis down")) + + await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index c5422e0d70f..9cd1fbb59a6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -373,3 +373,132 @@ class TestAnthropicMessagesHandlerToolInjection: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestAnthropicMessagesIncrementalScan: + """PR #33278: only_scan_new_messages through the real /v1/messages translation + handler (the path Claude Code uses). Encodes the wire payloads observed in the + live validation against a real Bedrock guardrail. + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + return BedrockGuardrail( + guardrail_name="bedrock-incremental-anthropic", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + def _data(self, messages, session_id): + return { + "model": "claude-sonnet-4-5", + "messages": messages, + "system": "You are a helpful geography assistant.", + "litellm_session_id": session_id, + } + + @pytest.mark.asyncio + async def test_first_turn_scans_all_eligible_then_second_turn_scans_only_diff(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-diff" + turn1 = [{"role": "user", "content": "What is the capital of France?"}] + turn2 = turn1 + [ + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data=self._data(turn1, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital of France?" + ] + mock_api.reset_mock() + await handler.process_input_messages( + data=self._data(turn2, sid), guardrail_to_apply=guardrail + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "Paris.", + "What is the capital of Germany?", + ] + + @pytest.mark.asyncio + async def test_identical_resend_makes_no_guardrail_call(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-resend" + msgs = [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "What is the capital of Germany?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.assert_not_called() + + @pytest.mark.asyncio + async def test_edited_history_message_is_rescanned(self): + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-edit" + msgs = [{"role": "user", "content": "What is the capital of France?"}] + edited = [{"role": "user", "content": "What is the capital and population of France?"}] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + mock_api.reset_mock() + await handler.process_input_messages(data=self._data(edited, sid), guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "What is the capital and population of France?" + ] + + @pytest.mark.asyncio + async def test_mixed_text_and_tool_use_keeps_text_segments(self): + """A message carrying both text and a tool_use block must not lose its text. + (tool_use inputs and tool_result content are dropped from texts on the + anthropic input path today; that is pre-existing baseline behavior.)""" + from unittest.mock import AsyncMock, patch + + handler = AnthropicMessagesHandler() + guardrail = self._bedrock_guardrail() + sid = "anth-sess-tools" + msgs = [ + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me look that up for you."}, + {"type": "tool_use", "id": "toolu_1", "name": "search", "input": {"query": "canary-args"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "canary-result"}], + }, + {"role": "user", "content": "Thanks, summarize the result."}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert "Let me look that up for you." in scanned, "text beside a tool_use must be scanned" + assert "Search for the weather in Paris" in scanned + assert "Thanks, summarize the result." in scanned diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4c268d9dfc9..7730b664c5e 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1137,3 +1137,95 @@ class TestGetStructuredMessages: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class TestIncrementalScanRespectsSkipFlags: + """PR #33278: skip_system_message_in_guardrail and skip_tool_message_in_guardrail + are enforced while this handler builds inputs["texts"] (_extract_inputs early + returns for system/tool roles), upstream of BedrockGuardrail's incremental path. + Bypassing _select_messages_for_apply_guardrail therefore cannot resurrect skipped + content on any turn, including a session's first turn where every segment is new. + Verified live against a real Bedrock ApplyGuardrail before being encoded here. + The flags are set as instance attributes, mirroring how guardrail_registry + applies litellm_params to the callback (they are not constructor kwargs). + """ + + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-incremental-skip-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + guardrail.skip_system_message_in_guardrail = True + guardrail.skip_tool_message_in_guardrail = True + return guardrail + + def _messages(self, followup=None): + base = [ + {"role": "system", "content": "SYSTEM-PROMPT-must-not-be-scanned"}, + {"role": "user", "content": "Search for the weather in Paris"}, + { + "role": "assistant", + "content": "Let me look that up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": '{"query": "weather"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-must-not-be-scanned"}, + {"role": "user", "content": "Thanks, summarize."}, + ] + return base + (followup or []) + + @pytest.mark.asyncio + async def test_first_turn_scans_no_system_or_tool_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = {"messages": self._messages(), "litellm_session_id": "skip-flags-turn1"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == [ + "Search for the weather in Paris", + "Let me look that up.", + "Thanks, summarize.", + ] + assert not any("SYSTEM-PROMPT" in text for text in scanned) + assert not any("TOOL-RESULT" in text for text in scanned) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_eligible_content(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + session = "skip-flags-turn2" + followup = [ + {"role": "assistant", "content": "It is sunny in Paris."}, + {"role": "user", "content": "And tomorrow?"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages( + data={"messages": self._messages(), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + mock_api.reset_mock() + await handler.process_input_messages( + data={"messages": self._messages(followup), "litellm_session_id": session}, + guardrail_to_apply=guardrail, + ) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["It is sunny in Paris.", "And tomorrow?"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 15827b80bcf..53f32fd96fb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3274,3 +3274,343 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n # CustomStreamWrapper would raise AttributeError inside __init__ and this # call would never reach here. assert response is not None + + +class TestBedrockOnlyScanNewMessages: + """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. + + apply_guardrail is the path the proxy actually runs for Bedrock (via the unified + guardrail interface), so these tests exercise it directly rather than the legacy + async_pre_call_hook. Each test uses a unique session id to isolate the process-wide + incremental cache. + """ + + def _guardrail(self): + return BedrockGuardrail( + guardrail_name="bedrock-incremental", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_messages(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-diff"} + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = bedrock_none + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + first_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in first_scanned] == ["be helpful", "first question"] + + mock_api.reset_mock() + + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question", "first answer", "second question"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 1 + second_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in second_scanned] == ["first answer", "second question"] + + @pytest.mark.asyncio + async def test_identical_resend_skips_api_call(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-resend"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + + mock_api.reset_mock() + result = await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" + ) + mock_api.assert_not_called() + assert result["texts"] == ["only question"] + + @pytest.mark.asyncio + async def test_no_session_id_scans_full_context(self): + guardrail = self._guardrail() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, + request_data={"metadata": {}}, + input_type="request", + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["q1", "a1", "q2"] + + @pytest.mark.asyncio + async def test_masking_guardrail_falls_back_and_does_not_persist(self): + """A guardrail that anonymizes content must not be short-circuited. + + Regression: the incremental fast path used to ignore the guardrail response, + so masked/anonymized output was dropped, the raw text reached the model, and + the segment was marked scanned so it was never re-checked. Detecting masked + output must force a full-context scan (which applies the masking) and must not + persist session state, so an identical resend is scanned again. + """ + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-mask"} + masked = { + "action": "GUARDRAIL_INTERVENED", + "output": [], + "outputs": [{"text": "my ssn is [REDACTED]"}], + } + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count == 2 + assert result["texts"] == ["my ssn is [REDACTED]"] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, + request_data=session, + input_type="request", + ) + assert mock_api.call_count >= 1 + first_scanned = mock_api.call_args_list[0].kwargs.get("messages") + assert first_scanned is not None + assert [m["content"] for m in first_scanned] == ["my ssn is 123-45-6789"] + + @pytest.mark.asyncio + async def test_generic_agent_multi_turn_scans_only_new_each_turn(self): + """A generic agent (not Claude Code) opts in by propagating a session id. + + Agent frameworks on the OpenAI SDK carry the session through the request + body (metadata.session_id here), not the x-claude-code-session-id header. + Across a growing multi-turn conversation every turn after the first must + send Bedrock only the newly appended segments, never the whole context. + """ + guardrail = self._guardrail() + session = {"metadata": {"session_id": "agent-multi-turn"}} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "system prompt", + "turn 1 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["system prompt", "turn 1 question", "turn 1 answer", "turn 2 question"]}, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 1 answer", + "turn 2 question", + ] + + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={ + "texts": [ + "system prompt", + "turn 1 question", + "turn 1 answer", + "turn 2 question", + "turn 2 answer", + "turn 3 question", + ] + }, + request_data=session, + input_type="request", + ) + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ + "turn 2 answer", + "turn 3 question", + ] + + def test_incremental_scan_cache_prefers_proxy_shared_cache(self): + guardrail = self._guardrail() + shared = DualCache() + proxy_logging = MagicMock() + proxy_logging.internal_usage_cache.dual_cache = shared + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging): + assert guardrail._incremental_scan_cache() is shared + + def test_incremental_scan_cache_falls_back_when_proxy_logging_missing(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch("litellm.proxy.proxy_server.proxy_logging_obj", None): + assert guardrail._incremental_scan_cache() is fallback_cache + + def test_incremental_scan_cache_falls_back_when_proxy_not_importable(self): + from litellm.integrations.custom_guardrail import dc as fallback_cache + + guardrail = self._guardrail() + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": None}): + assert guardrail._incremental_scan_cache() is fallback_cache + + @pytest.mark.asyncio + async def test_blocked_turn_is_rescanned_on_retry(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-blocked"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = HTTPException(status_code=400, detail="blocked") + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + + mock_api.reset_mock() + mock_api.side_effect = None + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["blocked prompt"] + + +class TestBedrockIncrementalFlagInteractions: + """Regression coverage for only_scan_new_messages combined with the other + Bedrock guardrail flags, from the PR #33278 live validation. Live evidence: + each of these was reproduced against a real Bedrock ApplyGuardrail first; + the mocks here encode the wire payloads observed there. + """ + + def _guardrail(self, **overrides): + params = dict( + guardrail_name="bedrock-incremental-flags", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + params.update(overrides) + return BedrockGuardrail(**params) + + @pytest.mark.asyncio + async def test_edited_history_segment_rescans_only_that_segment(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-edit"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, request_data=session, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1 EDITED", "a1", "q2"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["q1 EDITED"] + + @pytest.mark.asyncio + async def test_same_content_different_session_rescans_everything(self): + guardrail = self._guardrail() + texts = ["shared question", "shared answer"] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x1"}, input_type="request" + ) + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x2"}, input_type="request" + ) + assert mock_api.call_count == 1 + assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == texts + + @pytest.mark.asyncio + async def test_litellm_masking_flag_disables_incremental_single_full_scan(self): + """mask_request_content must fall back to exactly ONE full scan per turn + and never persist hashes (verified live: 1 call/turn, no cache writes).""" + guardrail = self._guardrail(mask_request_content=True) + session = {"litellm_session_id": "sess-flags-mask"} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" + + @pytest.mark.asyncio + async def test_server_side_anonymize_falls_back_full_scan_and_never_persists(self): + """A guardrail that rewrites content (Bedrock-side ANONYMIZE) must fall back + to the full scan so masking applies, and record no session state. Live + validation showed this costs 2 provider calls per turn; the count is + asserted here as documentation of that intended-tradeoff behavior.""" + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-flags-anon"} + masked = {"action": "NONE", "output": [{"text": "MASKED q1"}], "outputs": [{"text": "MASKED q1"}]} + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = masked + result = await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" + assert result["texts"] == ["MASKED q1"], "masked content must be applied" + mock_api.reset_mock() + await guardrail.apply_guardrail( + inputs={"texts": ["q1"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" + + @pytest.mark.asyncio + @pytest.mark.xfail( + reason="PR #33278 known gap: incremental path bypasses _select_messages_for_apply_guardrail, " + "so experimental_use_latest_role_message_only is silently ignored. Intended semantics " + "(pending DRI decision): incremental mode defers to the latest-role selection.", + strict=False, + ) + async def test_latest_role_only_is_respected_with_incremental(self): + guardrail = self._guardrail(experimental_use_latest_role_message_only=True) + session = {"litellm_session_id": "sess-flags-latestrole"} + structured = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q1"}, + ] + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["sys", "q1"], "structured_messages": structured}, + request_data=session, + input_type="request", + ) + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["q1"], "latest-role selection must exclude the system prompt" From 17a83aa89665ee5e640c9a032dd6d51b5b127cb5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:29:34 -0700 Subject: [PATCH 035/130] fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache (#33261) * fix(proxy): share CLI SSO login sessions across workers without enable_redis_auth_cache * fix(proxy): make CLI SSO flow state redis-authoritative across workers The CLI SSO flow is stored in a DualCache whose get_cache is memory-first, so the worker that served /sso/cli/start keeps serving its stale in-memory flow and never observes the sso_complete/session_data update another worker writes during the OAuth callback. Attaching Redis alone is not enough; poll on the original worker returns pending forever. Read and write the flow directly through the attached Redis backend when present so every worker sees the same authoritative state, falling back to the in-memory DualCache only when no Redis is configured. * fix(proxy): serialize CLI SSO flow as JSON for the redis round trip RedisCache stores values via str(value) and parses reads with json.loads then ast.literal_eval. The completed flow contains a LitellmUserRoles enum in session_data.user_role, whose repr is not a parseable literal, so any worker reading the completed flow from redis raised SyntaxError and returned 400 "CLI login session not found". Writing the flow as json.dumps makes the round trip lossless (the enum is a str subclass) and fails loudly at write time if a non-serializable value is ever added to the flow. * fix(proxy): point CLI SSO session-not-found hint at configuring Redis The error message and warning still told users to set enable_redis_auth_cache, but the CLI SSO session cache now gets Redis unconditionally whenever one is configured, so that flag no longer affects CLI login --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- litellm/proxy/management_endpoints/ui_sso.py | 56 +++-- litellm/proxy/proxy_server.py | 15 +- .../proxy/management_endpoints/test_ui_sso.py | 197 ++++++++++++++---- .../proxy/test_redis_auth_cache_flag.py | 40 +++- 4 files changed, 240 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3c8444ecf26..de988a0140f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -12,6 +12,7 @@ import asyncio import base64 import hashlib import inspect +import json import os import re import secrets @@ -258,11 +259,20 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic raise HTTPException(status_code=400, detail="Invalid CLI login session id") cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) - flow = cache.get_cache(key=cache_key) + redis_cache = cache.redis_cache + if redis_cache is not None: + flow = redis_cache.get_cache(key=cache_key) + else: + flow = cache.get_cache(key=cache_key) + if isinstance(flow, str): + try: + flow = json.loads(flow) + except ValueError: + flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: verbose_proxy_logger.warning( "CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, " - "a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.", + "a shared Redis cache is required for CLI login to work.", login_id, ) raise HTTPException( @@ -270,7 +280,7 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic detail=( "CLI login session not found or expired. Run `litellm-proxy login` again. " "If this happens immediately after starting a login, the proxy is likely running multiple " - "replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` " + "replicas without a shared cache; configure a Redis cache " "so every replica can see the login session." ), ) @@ -278,11 +288,12 @@ def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dic def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: - cache.set_cache( - key=_get_cli_sso_flow_cache_key(login_id), - value=flow, - ttl=CLI_SSO_SESSION_TTL_SECONDS, - ) + cache_key = _get_cli_sso_flow_cache_key(login_id) + redis_cache = cache.redis_cache + if redis_cache is not None: + redis_cache.set_cache(key=cache_key, value=json.dumps(flow), ttl=CLI_SSO_SESSION_TTL_SECONDS) + else: + cache.set_cache(key=cache_key, value=flow, ttl=CLI_SSO_SESSION_TTL_SECONDS) def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: @@ -593,11 +604,11 @@ def _render_cli_sso_verification_page( @router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False) async def cli_sso_start(request: Request): - from litellm.proxy.proxy_server import general_settings, user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings _check_cli_sso_start_rate_limit( request=request, - cache=user_api_key_cache, + cache=cli_sso_session_cache, use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), ) @@ -612,7 +623,7 @@ async def cli_sso_start(request: Request): "user_code_verified": False, "session_data": None, } - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) verification_uri_complete: str | None = ( ( @@ -644,9 +655,9 @@ async def cli_sso_complete(request: Request, login_id: str): from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache - flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache) if not flow.get("sso_complete") or not flow.get("session_data"): raise HTTPException(status_code=400, detail="CLI login is not ready") @@ -670,7 +681,7 @@ async def cli_sso_complete(request: Request, login_id: str): raise HTTPException(status_code=400, detail="Invalid verification code") flow["user_code_verified"] = True - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) @@ -861,10 +872,10 @@ async def google_login( Example: """ from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, premium_user, prisma_client, - user_api_key_cache, user_custom_ui_sso_sign_in_handler, ) @@ -912,7 +923,7 @@ async def google_login( ) if source == LITELLM_CLI_SOURCE_IDENTIFIER: - _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) # Store CLI login handle in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( @@ -1957,6 +1968,7 @@ async def _complete_cli_sso_callback_session( user_defined_values: Optional[SSOUserDefinedValues], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, + cli_sso_session_cache: DualCache, proxy_logging_obj: ProxyLogging, prefill_user_code: str | None = None, sso_assertion: SSOIdentityAssertion | None = None, @@ -2006,7 +2018,7 @@ async def _complete_cli_sso_callback_session( flow["sso_complete"] = True browser_complete_token = secrets.token_urlsafe(32) flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token) - _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow) + _set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow) verbose_proxy_logger.info( f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" @@ -2037,13 +2049,14 @@ async def cli_sso_callback( verbose_proxy_logger.info("CLI SSO callback") from litellm.proxy.proxy_server import ( + cli_sso_session_cache, general_settings, prisma_client, proxy_logging_obj, user_api_key_cache, ) - flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key, cache=cli_sso_session_cache) if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) @@ -2083,6 +2096,7 @@ async def cli_sso_callback( user_defined_values=user_defined_values, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + cli_sso_session_cache=cli_sso_session_cache, proxy_logging_obj=proxy_logging_obj, prefill_user_code=prefill_user_code, sso_assertion=sso_assertion, @@ -2114,10 +2128,10 @@ async def cli_poll_key( team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. """ from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import cli_sso_session_cache try: - flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) + flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=cli_sso_session_cache) if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): raise HTTPException(status_code=403, detail="Invalid CLI polling secret") @@ -2192,7 +2206,7 @@ async def cli_poll_key( ) # Delete cache entry (single-use) - user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) + cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}") poll_response = { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6de3e43fc1a..d4bff81ea6a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -226,6 +226,7 @@ from litellm.constants import ( APSCHEDULER_MAX_INSTANCES, APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, + CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -1970,6 +1971,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) +cli_sso_session_cache = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits @@ -3696,13 +3698,22 @@ def _build_redis_usage_cache_from_environment() -> RedisCache | None: def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: bool) -> None: """ Wires an established coordination Redis into the proxy-level caches that - consume it directly: the spend counter cache, the cluster-wide config - cache, and (only when opted in) the virtual-key auth cache. + consume it directly: the spend counter cache, the CLI SSO login-session + cache, the cluster-wide config cache, and (only when opted in) the + virtual-key auth cache. + + The CLI SSO login-session cache is always backed by Redis when available so + that the browser SSO flow behind `lite login` survives landing on different + workers; it must not be gated behind enable_redis_auth_cache. """ spend_counter_cache.attach_redis_cache( redis_cache, default_redis_ttl=litellm.default_redis_ttl, ) + cli_sso_session_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, + ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e1856860c8a..47ceb0c05fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2214,7 +2214,95 @@ class TestCLIKeyRegenerationFlow: _get_cli_sso_flow_or_raise(login_id="cli-test_1234567890", cache=mock_cache) assert expired_exc.value.status_code == 400 assert "session not found or expired" in expired_exc.value.detail - assert "enable_redis_auth_cache" in expired_exc.value.detail + assert "configure a Redis cache" in expired_exc.value.detail + assert "enable_redis_auth_cache" not in expired_exc.value.detail + + def test_cli_sso_flow_is_redis_authoritative_when_redis_attached(self): + """ + When Redis is attached, the CLI SSO flow must be read from and written to + Redis directly, never the in-memory layer. Otherwise the worker that served + /sso/cli/start keeps serving its stale in-memory flow and never sees the + sso_complete/session_data update another worker wrote, which is exactly the + multi-worker failure this fix targets. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + CLI_SSO_SESSION_TTL_SECONDS, + _get_cli_sso_flow_cache_key, + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-redis_authoritative_1234567890" + cache_key = _get_cli_sso_flow_cache_key(login_id) + fresh_flow = {"poll_secret_hash": "fresh", "sso_complete": True} + stale_flow = {"poll_secret_hash": "stale", "sso_complete": False} + + redis_cache = MagicMock() + redis_cache.get_cache.return_value = fresh_flow + cache = MagicMock() + cache.redis_cache = redis_cache + cache.get_cache.return_value = stale_flow + + result = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert result == fresh_flow + redis_cache.get_cache.assert_called_once_with(key=cache_key) + cache.get_cache.assert_not_called() + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=fresh_flow) + + redis_cache.set_cache.assert_called_once_with( + key=cache_key, value=json.dumps(fresh_flow), ttl=CLI_SSO_SESSION_TTL_SECONDS + ) + cache.set_cache.assert_not_called() + + def test_cli_sso_flow_with_enum_survives_redis_round_trip(self): + """ + RedisCache stores values via str(value) and reads them back through + json.loads/ast.literal_eval. A raw flow dict containing a Python enum + (session_data.user_role after the SSO callback) produces an unparseable + repr, so every worker reading the completed flow from Redis got a + SyntaxError and returned 400 "session not found". The flow must survive + a real Redis serialization round trip. + """ + from litellm.caching.redis_cache import RedisCache + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_or_raise, + _set_cli_sso_flow, + ) + + login_id = "cli-enum_round_trip_1234567890" + completed_flow = { + "poll_secret_hash": "hash", + "sso_complete": True, + "user_code_verified": False, + "session_data": { + "user_id": "user-1", + "user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + "models": [], + "teams": ["team-1"], + "team_details": [{"team_id": "team-1", "team_alias": "alias"}], + }, + } + + redis_store: dict = {} + redis_cache = MagicMock() + redis_cache.set_cache.side_effect = lambda key, value, ttl: redis_store.__setitem__( + key, str(value).encode("utf-8") + ) + redis_cache.get_cache.side_effect = lambda key: RedisCache._get_cache_logic( + MagicMock(), redis_store.get(key) + ) + cache = MagicMock() + cache.redis_cache = redis_cache + + _set_cli_sso_flow(login_id=login_id, cache=cache, flow=completed_flow) + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cache) + + assert flow["sso_complete"] is True + assert flow["session_data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + assert flow["session_data"]["team_details"] == [{"team_id": "team-1", "team_alias": "alias"}] @pytest.mark.asyncio async def test_cli_sso_start_creates_bound_flow(self): @@ -2228,10 +2316,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_sso_start(request=mock_request) assert result["login_id"].startswith("cli-") @@ -2259,10 +2350,13 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 31 - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_start(request=mock_request) @@ -2281,7 +2375,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2315,7 +2409,7 @@ class TestCLIKeyRegenerationFlow: mock_request.client = SimpleNamespace(host="127.0.0.1") mock_request.headers = {} mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.increment_cache.return_value = 1 with ( @@ -2349,7 +2443,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.example.com/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} async def drive(enabled: bool): @@ -2358,6 +2452,7 @@ class TestCLIKeyRegenerationFlow: patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", None, @@ -2525,7 +2620,7 @@ class TestCLIKeyRegenerationFlow: ) mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2544,6 +2639,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), ): result = await cli_sso_callback( request=mock_request, @@ -2568,7 +2664,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2582,6 +2678,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2606,7 +2703,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH") - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2618,7 +2715,10 @@ class TestCLIKeyRegenerationFlow: "session_data": {"user_id": "test-user-123"}, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2640,7 +2740,7 @@ class TestCLIKeyRegenerationFlow: mock_request.body = AsyncMock( return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "user_code_hash": _hash_cli_sso_secret( @@ -2651,7 +2751,10 @@ class TestCLIKeyRegenerationFlow: "session_data": None, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_sso_complete( request=mock_request, login_id="cli-session-4567890" @@ -2687,7 +2790,7 @@ class TestCLIKeyRegenerationFlow: mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -2709,6 +2812,7 @@ class TestCLIKeyRegenerationFlow: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", @@ -2769,7 +2873,7 @@ class TestCLIKeyRegenerationFlow: } # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2777,7 +2881,10 @@ class TestCLIKeyRegenerationFlow: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): # Act - First poll without team_id result = await cli_poll_key( key_id=session_key, @@ -2803,7 +2910,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2816,7 +2923,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): with pytest.raises(HTTPException) as exc_info: await cli_poll_key(key_id="cli-session-789123", team_id=None) @@ -2830,7 +2940,7 @@ class TestCLIKeyRegenerationFlow: cli_poll_key, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -2843,7 +2953,10 @@ class TestCLIKeyRegenerationFlow: }, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id="cli-session-789123", team_id=None, @@ -3011,7 +3124,7 @@ class TestCLIKeyRegenerationFlow: ) # Mock cache - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3023,6 +3136,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3086,7 +3200,7 @@ class TestCLIKeyRegenerationFlow: models=["gpt-4"], max_budget=100.0, ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3097,6 +3211,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", @@ -3142,7 +3257,7 @@ class TestCLIKeyRegenerationFlow: "models": ["gpt-4"], "user_email": "unbudgeted@example.com", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -3153,6 +3268,7 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", return_value=mock_jwt_token, @@ -4082,7 +4198,7 @@ class TestPKCEFunctionality: mock_request.query_params = {"state": test_state} # Mock cache with async methods — use dict format (primary path) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) test_code_verifier = "test_code_verifier_abc123xyz" mock_cache.async_get_cache = AsyncMock( return_value={"code_verifier": test_code_verifier} @@ -4133,7 +4249,7 @@ class TestPKCEFunctionality: mock_sso.__exit__ = MagicMock(return_value=False) test_state = "test456" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_set_cache = AsyncMock() @@ -4657,7 +4773,7 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4783,7 +4899,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4825,7 +4941,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=None) # verifier not found mock_request = MagicMock(spec=Request) @@ -4913,7 +5029,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=12345) mock_cache.async_delete_cache = AsyncMock() @@ -4965,7 +5081,7 @@ class TestPKCEFunctionality: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler legacy_verifier = "legacy_plain_string_verifier_abc123" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier) mock_request = MagicMock(spec=Request) @@ -6249,7 +6365,7 @@ class TestCliSsoAttributionMetadata: provider="generic", team_ids=[], ) - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6266,6 +6382,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), ): await ui_sso.cli_sso_callback( @@ -6290,7 +6407,7 @@ class TestCliSsoAttributionMetadata: mock_request = MagicMock(spec=Request) mock_request.base_url = "http://internal-proxy.local/" - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6313,6 +6430,7 @@ class TestCliSsoAttributionMetadata: ) as get_user_info_mock, patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.proxy_server.general_settings", @@ -6359,7 +6477,7 @@ class TestCliSsoAttributionMetadata: "user_id": "test-user-123", "employment_type": "contractor", } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": "poll-secret-hash", "user_code_hash": "user-code-hash", @@ -6387,6 +6505,7 @@ class TestCliSsoAttributionMetadata: ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.user_custom_sso", None), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", @@ -6428,7 +6547,7 @@ class TestCliSsoAttributionMetadata: "org": {"cost_center": "CC-42"}, }, } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -6436,7 +6555,10 @@ class TestCliSsoAttributionMetadata: "session_data": session_data, } - with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + ): result = await cli_poll_key( key_id=session_key, team_id=None, @@ -7287,7 +7409,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): "models": ["gpt-4"], } - mock_cache = MagicMock() + mock_cache = MagicMock(redis_cache=None) mock_cache.get_cache.return_value = { "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), "sso_complete": True, @@ -7299,6 +7421,7 @@ async def test_cli_poll_key_tolerates_missing_user_row(): with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index d0cb5ec5465..849d5494c6e 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -54,8 +54,8 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): _FakeRedisCache (passes the isinstance guard in _init_cache). 3. Extracts enable_redis_auth_cache from litellm_settings and passes it as the second argument to _init_cache (matching production behaviour). - 4. Yields (user_api_key_cache, spend_counter_cache) after calling - _init_cache, then restores everything. + 4. Yields (user_api_key_cache, spend_counter_cache, cli_sso_session_cache) + after calling _init_cache, then restores everything. """ fake_redis = _FakeRedisCache() @@ -64,19 +64,21 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): fresh_user_cache = DualCache() fresh_spend_cache = DualCache() + fresh_cli_sso_cache = DualCache() enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False) with ( patch.object(ps, "user_api_key_cache", fresh_user_cache), patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. patch("litellm.Cache", return_value=mock_litellm_cache), ): litellm.cache = None ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache) - yield fresh_user_cache, fresh_spend_cache + yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache # --------------------------------------------------------------------------- @@ -90,7 +92,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": True}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is not None, ( "Redis should be attached to user_api_key_cache when " "enable_redis_auth_cache=True" @@ -101,7 +103,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache=False" @@ -112,7 +114,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, _): + ) as (user_cache, _, _cli_sso_cache): assert user_cache.redis_cache is None, ( "user_api_key_cache must remain in-memory-only when " "enable_redis_auth_cache is absent from litellm_settings" @@ -129,7 +131,7 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings=ls, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (_, spend_cache): + ) as (_, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None, ( f"spend_counter_cache must always get Redis " f"(enable_redis_auth_cache={flag_value!r})" @@ -140,6 +142,28 @@ class TestRedisAuthCacheFlag: with _patched_init_cache( litellm_settings={"enable_redis_auth_cache": False}, cache_params={"type": "redis", "host": "localhost", "port": 6379}, - ) as (user_cache, spend_cache): + ) as (user_cache, spend_cache, _cli_sso_cache): assert spend_cache.redis_cache is not None assert user_cache.redis_cache is None + + def test_cli_sso_session_cache_always_gets_redis_regardless_of_flag(self): + """ + cli_sso_session_cache must receive Redis regardless of the auth-cache + flag so that `lite login` works on multi-worker deployments without + enable_redis_auth_cache (regression for the CLI SSO "Invalid CLI login + session" bug) + """ + for flag_value in (True, False, None): + ls = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + litellm_settings=ls, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, _, cli_sso_cache): + assert cli_sso_cache.redis_cache is not None, ( + f"cli_sso_session_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) From 0fcaadf11ca1676f5f3e041808caf837fdb71ee3 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 22 Jul 2026 12:43:10 -0700 Subject: [PATCH 036/130] test(e2e): move Admin UI Playwright suite to tests/e2e/ui (#34196) Relocates ui/litellm-dashboard/e2e_tests to tests/e2e/ui so all end to end suites live under tests/e2e. The suite stays in TypeScript and becomes a self-contained npm package with its own package.json, lockfile and tsconfig instead of leaning on the dashboard's toolchain; the dashboard drops its @playwright/test dependency, e2e scripts and knip/vitest/tsconfig carve-outs. CI paths follow the move: both CircleCI jobs (main e2e and the SERVER_ROOT_PATH migration smoke) and the test_server_root_path workflow now install and run Playwright from tests/e2e/ui, with the node cache keyed on both lockfiles. classify_changes.sh treats tests/e2e/ui as client so spec edits keep skipping backend jobs. The suite's mock LLM fixture is excluded from the e2e basedpyright zero-error gate in pyrightconfig.json since it belongs to the TS suite, not the typed Python harness. --- .circleci/config.yml | 42 ++++--- .circleci/scripts/classify_changes.sh | 2 +- .github/workflows/test_server_root_path.yml | 10 +- pyrightconfig.json | 2 +- tests/e2e/CLAUDE.md | 1 + tests/e2e/load/test_session_anomaly.py | 4 +- .../e2e_tests => tests/e2e/ui}/constants.ts | 0 .../e2e/ui}/fixtures/config.yml | 0 .../e2e/ui}/fixtures/menuMappings.ts | 0 .../e2e/ui}/fixtures/migratedPages.ts | 0 .../ui}/fixtures/mock_llm_server/server.py | 0 .../e2e/ui}/fixtures/pages.ts | 0 .../e2e/ui}/fixtures/roles.ts | 0 .../e2e/ui}/fixtures/seed.sql | 0 .../e2e/ui}/fixtures/users.ts | 0 .../e2e_tests => tests/e2e/ui}/globalSetup.ts | 0 .../e2e/ui}/helpers/navigation.ts | 0 .../ui}/migration.serverRootPath.config.ts | 0 .../migration.serverRootPath.globalSetup.ts | 0 tests/e2e/ui/package-lock.json | 111 ++++++++++++++++++ tests/e2e/ui/package.json | 16 +++ .../e2e/ui}/playwright.config.ts | 0 .../e2e_tests => tests/e2e/ui}/run_e2e.sh | 6 +- .../e2e/ui}/serverRootPath.config.ts | 0 .../e2e/ui}/tests/auth/logout.spec.ts | 0 .../e2e/ui}/tests/auth/proxyLogoutUrl.spec.ts | 0 .../auth/unauthenticatedRedirect.spec.ts | 0 .../tests/internal-user/internalUser.spec.ts | 0 .../internal-user/internalUserNoTeam.spec.ts | 0 .../internalUserWithTeams.spec.ts | 0 .../internal-viewer/internalViewer.spec.ts | 0 .../tests/login/internalUserIdentity.spec.ts | 0 .../e2e/ui}/tests/login/login.spec.ts | 0 .../login/serverRootPathRedirect.spec.ts | 0 .../e2e/ui}/tests/mcp/mcpServers.spec.ts | 0 .../e2e/ui}/tests/migration/README.md | 5 +- .../ui}/tests/migration/migratedPages.spec.ts | 0 .../e2e/ui}/tests/modelHub/modelHub.spec.ts | 0 .../e2e/ui}/tests/modelsPage/addModel.spec.ts | 0 .../modelsPage/clearCustomPricing.spec.ts | 0 .../ui}/tests/modelsPage/credentials.spec.ts | 0 .../e2e/ui}/tests/navigation/sidebar.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/keys.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/license.spec.ts | 0 .../e2e/ui}/tests/proxy-admin/teams.spec.ts | 0 .../ui}/tests/settings/adminSettings.spec.ts | 0 .../ui}/tests/settings/routerSettings.spec.ts | 2 +- .../ui}/tests/team-admin/teamAdmin.spec.ts | 0 .../e2e/ui}/tests/users/searchUsers.spec.ts | 0 .../ui}/tests/users/viewInternalUsers.spec.ts | 0 tests/e2e/ui/tsconfig.json | 16 +++ .../proxy/management_endpoints/test_ui_sso.py | 1 + ui/litellm-dashboard/knip.json | 10 +- ui/litellm-dashboard/package-lock.json | 11 +- ui/litellm-dashboard/package.json | 5 - ui/litellm-dashboard/tsconfig.json | 2 +- ui/litellm-dashboard/vitest.config.ts | 3 +- 57 files changed, 194 insertions(+), 55 deletions(-) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/constants.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/config.yml (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/menuMappings.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/migratedPages.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/mock_llm_server/server.py (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/pages.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/roles.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/seed.sql (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/fixtures/users.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/globalSetup.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/helpers/navigation.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/migration.serverRootPath.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/migration.serverRootPath.globalSetup.ts (100%) create mode 100644 tests/e2e/ui/package-lock.json create mode 100644 tests/e2e/ui/package.json rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/playwright.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/run_e2e.sh (98%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/serverRootPath.config.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/logout.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/proxyLogoutUrl.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/auth/unauthenticatedRedirect.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUser.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUserNoTeam.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-user/internalUserWithTeams.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/internal-viewer/internalViewer.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/internalUserIdentity.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/login.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/login/serverRootPathRedirect.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/mcp/mcpServers.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/migration/README.md (84%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/migration/migratedPages.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelHub/modelHub.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/addModel.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/clearCustomPricing.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/modelsPage/credentials.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/navigation/sidebar.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/keys.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/license.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/proxy-admin/teams.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/settings/adminSettings.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/settings/routerSettings.spec.ts (98%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/team-admin/teamAdmin.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/users/searchUsers.spec.ts (100%) rename {ui/litellm-dashboard/e2e_tests => tests/e2e/ui}/tests/users/viewInternalUsers.spec.ts (100%) create mode 100644 tests/e2e/ui/tsconfig.json diff --git a/.circleci/config.yml b/.circleci/config.yml index b0a705966a2..2f01b6de4f3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2731,7 +2731,7 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright # The cimg/python:3.12-browsers image already ships the Chromium system @@ -2742,11 +2742,14 @@ jobs: command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2777,10 +2780,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy @@ -2798,7 +2801,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2819,15 +2822,15 @@ jobs: # Forward LITELLM_LICENSE so license.spec.ts can detect that the # proxy was launched with a license and assert premium_user=true. command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/playwright.config.ts + npx playwright test --config playwright.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-playwright-report e2e_ui_testing_server_root_path: @@ -2870,17 +2873,20 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright command: | cd ui/litellm-dashboard npm ci + cd ../../tests/e2e/ui + npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules + - tests/e2e/ui/node_modules - ~/.cache/ms-playwright - run: name: Build UI from source @@ -2902,10 +2908,10 @@ jobs: name: Seed database command: | PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ - -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + -f tests/e2e/ui/fixtures/seed.sql - run: name: Start mock LLM server - command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true - run: name: Start LiteLLM proxy under a server root path @@ -2918,7 +2924,7 @@ jobs: command: | LITELLM_LICENSE="$LITELLM_LICENSE" \ uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --config tests/e2e/ui/fixtures/config.yml \ --port 4000 background: true - run: @@ -2937,15 +2943,15 @@ jobs: - run: name: Run migration smoke under SERVER_ROOT_PATH command: | - cd ui/litellm-dashboard + cd tests/e2e/ui LITELLM_LICENSE="$LITELLM_LICENSE" \ - npx playwright test --config e2e_tests/migration.serverRootPath.config.ts + npx playwright test --config migration.serverRootPath.config.ts no_output_timeout: 10m - store_artifacts: - path: ui/litellm-dashboard/test-results + path: tests/e2e/ui/test-results destination: e2e-server-root-path-test-results - store_artifacts: - path: ui/litellm-dashboard/playwright-report + path: tests/e2e/ui/playwright-report destination: e2e-server-root-path-playwright-report build_docker_database_image: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 2c15428be6a..2ca2654a207 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -8,7 +8,7 @@ has_backend=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in - ui/*) has_client=true ;; + ui/* | tests/e2e/ui/*) has_client=true ;; docs/* | *.md | *.mdx) : ;; *) has_backend=true ;; esac diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index f59cee29893..01f70511e79 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -106,8 +106,8 @@ jobs: with: node-version: "20" - - name: Install UI deps and Chromium - working-directory: ui/litellm-dashboard + - name: Install e2e deps and Chromium + working-directory: tests/e2e/ui run: | retry() { local attempt=1 @@ -131,17 +131,17 @@ jobs: retry npx playwright install --with-deps chromium - name: Run SERVER_ROOT_PATH redirect e2e - working-directory: ui/litellm-dashboard + working-directory: tests/e2e/ui env: SERVER_ROOT_PATH: ${{ matrix.root_path }} - run: npx playwright test --config=e2e_tests/serverRootPath.config.ts + run: npx playwright test --config=serverRootPath.config.ts - name: Upload Playwright artifacts on failure if: failure() uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 with: name: playwright-trace-${{ strategy.job-index }} - path: ui/litellm-dashboard/test-results/ + path: tests/e2e/ui/test-results/ retention-days: 7 - name: Cleanup diff --git a/pyrightconfig.json b/pyrightconfig.json index eabfbf515c4..2686ccd73d9 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,7 +1,7 @@ { "include": ["litellm"], "ignore": [], - "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "litellm/types/utils.py", "litellm/proxy/_types.py"], + "exclude": ["**/node_modules", "**/__pycache__", "tests/e2e/claude_code", "tests/e2e/ui", "litellm/types/utils.py", "litellm/proxy/_types.py"], "pythonVersion": "3.12", "typeCheckingMode": "strict", "enableTypeIgnoreComments": false, diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 186517290ea..0e39664e358 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -22,6 +22,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` ## MCP suite: real Datadog only diff --git a/tests/e2e/load/test_session_anomaly.py b/tests/e2e/load/test_session_anomaly.py index 80f8aff3ba4..7062587352b 100644 --- a/tests/e2e/load/test_session_anomaly.py +++ b/tests/e2e/load/test_session_anomaly.py @@ -62,7 +62,7 @@ class TestSummarizePlannedTurns: class TestRetried: def test_transient_failures_then_success_returns_the_success(self) -> None: - outcome = Success(data=SessionMessagesResponse()) + outcome = Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()) calls = iter( (NetworkError(message="overloaded"), NetworkError(message="overloaded"), outcome) ) @@ -88,7 +88,7 @@ class TestRetried: raise AssertionError("slept after a successful attempt") result = retried( - lambda: Success(data=SessionMessagesResponse()), + lambda: Success[SessionMessagesResponse](status_code=200, data=SessionMessagesResponse()), attempts=3, sleep=sleep_means_retry, ) diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/tests/e2e/ui/constants.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/constants.ts rename to tests/e2e/ui/constants.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/config.yml b/tests/e2e/ui/fixtures/config.yml similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/config.yml rename to tests/e2e/ui/fixtures/config.yml diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts rename to tests/e2e/ui/fixtures/menuMappings.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts rename to tests/e2e/ui/fixtures/migratedPages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py b/tests/e2e/ui/fixtures/mock_llm_server/server.py similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py rename to tests/e2e/ui/fixtures/mock_llm_server/server.py diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/pages.ts rename to tests/e2e/ui/fixtures/pages.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/tests/e2e/ui/fixtures/roles.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/roles.ts rename to tests/e2e/ui/fixtures/roles.ts diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/seed.sql rename to tests/e2e/ui/fixtures/seed.sql diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/tests/e2e/ui/fixtures/users.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/fixtures/users.ts rename to tests/e2e/ui/fixtures/users.ts diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/tests/e2e/ui/globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/globalSetup.ts rename to tests/e2e/ui/globalSetup.ts diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/helpers/navigation.ts rename to tests/e2e/ui/helpers/navigation.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts b/tests/e2e/ui/migration.serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts rename to tests/e2e/ui/migration.serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts b/tests/e2e/ui/migration.serverRootPath.globalSetup.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/migration.serverRootPath.globalSetup.ts rename to tests/e2e/ui/migration.serverRootPath.globalSetup.ts diff --git a/tests/e2e/ui/package-lock.json b/tests/e2e/ui/package-lock.json new file mode 100644 index 00000000000..b22673a3535 --- /dev/null +++ b/tests/e2e/ui/package-lock.json @@ -0,0 +1,111 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "litellm-ui-e2e", + "version": "0.0.0", + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } + }, + "node_modules/@playwright/test": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", + "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/@types/node": { + "version": "20.19.37", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.37.tgz", + "integrity": "sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, + "node_modules/playwright": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", + "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.58.1" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.58.1", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", + "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/tests/e2e/ui/package.json b/tests/e2e/ui/package.json new file mode 100644 index 00000000000..ede759d97cb --- /dev/null +++ b/tests/e2e/ui/package.json @@ -0,0 +1,16 @@ +{ + "name": "litellm-ui-e2e", + "version": "0.0.0", + "private": true, + "scripts": { + "e2e": "playwright test --config playwright.config.ts", + "e2e:ui": "playwright test --ui --config playwright.config.ts", + "e2e:migration": "playwright test tests/migration/migratedPages.spec.ts --config playwright.config.ts", + "e2e:migration:root": "playwright test --config migration.serverRootPath.config.ts" + }, + "devDependencies": { + "@playwright/test": "1.58.1", + "@types/node": "20.19.37", + "typescript": "5.9.3" + } +} diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/tests/e2e/ui/playwright.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/playwright.config.ts rename to tests/e2e/ui/playwright.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/run_e2e.sh b/tests/e2e/ui/run_e2e.sh similarity index 98% rename from ui/litellm-dashboard/e2e_tests/run_e2e.sh rename to tests/e2e/ui/run_e2e.sh index ea95f18890c..858eb401c8e 100755 --- a/ui/litellm-dashboard/e2e_tests/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -20,8 +20,8 @@ set -euo pipefail # ================================================================ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DASHBOARD_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" +DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard" IS_CI="${CI:-false}" CONTAINER_NAME="litellm-e2e-postgres-$$" MOCK_PID="" @@ -187,12 +187,12 @@ PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAM # --- Playwright --- echo "=== Installing Playwright dependencies ===" -cd "$DASHBOARD_DIR" +cd "$SCRIPT_DIR" npm install --silent 2>/dev/null || true npx playwright install chromium --with-deps 2>/dev/null || npx playwright install chromium echo "=== Running Playwright tests ===" -npx playwright test --config e2e_tests/playwright.config.ts "$@" +npx playwright test --config playwright.config.ts "$@" EXIT_CODE=$? exit $EXIT_CODE diff --git a/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts b/tests/e2e/ui/serverRootPath.config.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts rename to tests/e2e/ui/serverRootPath.config.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts rename to tests/e2e/ui/tests/auth/logout.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts b/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts rename to tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts rename to tests/e2e/ui/tests/auth/unauthenticatedRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUser.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUser.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserWithTeams.spec.ts rename to tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/internal-viewer/internalViewer.spec.ts rename to tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/tests/e2e/ui/tests/login/internalUserIdentity.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts rename to tests/e2e/ui/tests/login/internalUserIdentity.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/tests/e2e/ui/tests/login/login.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts rename to tests/e2e/ui/tests/login/login.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts b/tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts rename to tests/e2e/ui/tests/login/serverRootPathRedirect.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts b/tests/e2e/ui/tests/mcp/mcpServers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts rename to tests/e2e/ui/tests/mcp/mcpServers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md similarity index 84% rename from ui/litellm-dashboard/e2e_tests/tests/migration/README.md rename to tests/e2e/ui/tests/migration/README.md index 4b3a391d421..d6b33598ec4 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -9,8 +9,9 @@ the default mount and a non-root `SERVER_ROOT_PATH` mount. ## Adding a page When a page's migration merges, add its route segment to -`e2e_tests/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `src/utils/migratedPages.ts`). Both suites pick it up automatically. +`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` +in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up +automatically. ## Running diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts rename to tests/e2e/ui/tests/migration/migratedPages.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts rename to tests/e2e/ui/tests/modelHub/modelHub.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts rename to tests/e2e/ui/tests/modelsPage/addModel.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts rename to tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts rename to tests/e2e/ui/tests/modelsPage/credentials.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts rename to tests/e2e/ui/tests/navigation/sidebar.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts rename to tests/e2e/ui/tests/proxy-admin/keys.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts b/tests/e2e/ui/tests/proxy-admin/license.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts rename to tests/e2e/ui/tests/proxy-admin/license.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts rename to tests/e2e/ui/tests/proxy-admin/teams.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts b/tests/e2e/ui/tests/settings/adminSettings.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts rename to tests/e2e/ui/tests/settings/adminSettings.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts b/tests/e2e/ui/tests/settings/routerSettings.spec.ts similarity index 98% rename from ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts rename to tests/e2e/ui/tests/settings/routerSettings.spec.ts index 3e140b9ab56..ffa5f2c2ae2 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts +++ b/tests/e2e/ui/tests/settings/routerSettings.spec.ts @@ -6,7 +6,7 @@ import { Role, users } from "../../fixtures/users"; // Type-only import of the OpenAPI-generated backend schema, erased at runtime by // esbuild. It types the round-trips below so mistakes surface in the editor; the live // test against the real proxy is what actually enforces the contract. -import type { components } from "../../../src/lib/http/schema"; +import type { components } from "../../../../../ui/litellm-dashboard/src/lib/http/schema"; // These tests mutate the proxy's shared router_settings, and the Loadbalancing save // echoes the whole settings object, so they must not run concurrently. diff --git a/ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/team-admin/teamAdmin.spec.ts rename to tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts rename to tests/e2e/ui/tests/users/searchUsers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts similarity index 100% rename from ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts rename to tests/e2e/ui/tests/users/viewInternalUsers.spec.ts diff --git a/tests/e2e/ui/tsconfig.json b/tests/e2e/ui/tsconfig.json new file mode 100644 index 00000000000..f9290fe7b49 --- /dev/null +++ b/tests/e2e/ui/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "commonjs", + "moduleResolution": "node", + "lib": ["ES2022", "DOM"], + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "resolveJsonModule": true, + "noEmit": true, + "types": ["node"] + }, + "include": ["**/*.ts"], + "exclude": ["node_modules"] +} diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 47ceb0c05fa..c693017e134 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7756,6 +7756,7 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): user_defined_values=None, prisma_client=MagicMock(), user_api_key_cache=MagicMock(), + cli_sso_session_cache=MagicMock(), proxy_logging_obj=MagicMock(), sso_assertion=assertion, ) diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index afed6b0f90e..48b39e8122d 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,7 +1,7 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", "entry": ["scripts/**/*.{ts,mjs}", "src/components/ui/**/*.{ts,tsx}"], - "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}", "e2e_tests/**/*.ts"], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}"], "ignore": ["src/lib/http/schema.d.ts"], "ignoreDependencies": [ "openapi-typescript", @@ -10,14 +10,6 @@ "tailwindcss", "tw-animate-css" ], - "playwright": { - "config": [ - "e2e_tests/playwright.config.ts", - "e2e_tests/serverRootPath.config.ts", - "e2e_tests/migration.serverRootPath.config.ts" - ], - "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] - }, "vitest": { "config": ["vitest.config.ts"] }, diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 742c1e4a63f..14c8f0fdc18 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -47,7 +47,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", @@ -2723,8 +2722,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright": "1.58.1" }, @@ -7361,7 +7361,6 @@ "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -10948,8 +10947,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "dependencies": { "playwright-core": "1.58.1" }, @@ -10967,8 +10967,9 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", - "devOptional": true, "license": "Apache-2.0", + "optional": true, + "peer": true, "bin": { "playwright-core": "cli.js" }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 0f54c536297..5add2ad4e9e 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -14,10 +14,6 @@ "test:coverage": "vitest run --coverage", "format": "prettier --write .", "format:check": "prettier --check .", - "e2e": "playwright test --config e2e_tests/playwright.config.ts", - "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts", - "e2e:migration": "playwright test e2e_tests/tests/migration/migratedPages.spec.ts --config e2e_tests/playwright.config.ts", - "e2e:migration:root": "playwright test --config e2e_tests/migration.serverRootPath.config.ts", "knip": "knip", "knip:ci": "knip --exclude exports,nsExports,types,nsTypes,enumMembers,classMembers,duplicates", "knip:fix": "knip --fix", @@ -63,7 +59,6 @@ }, "devDependencies": { "@eslint/js": "9.39.2", - "@playwright/test": "1.58.1", "@tailwindcss/forms": "0.5.11", "@tailwindcss/postcss": "4.3.2", "@testing-library/dom": "10.4.1", diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json index 8ca1752013a..5ca97e3e9db 100644 --- a/ui/litellm-dashboard/tsconfig.json +++ b/ui/litellm-dashboard/tsconfig.json @@ -23,5 +23,5 @@ "target": "ES2017" }, "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts", ".next/dev/types/**/*.ts"], - "exclude": ["node_modules", "e2e_tests", "scripts"] + "exclude": ["node_modules", "scripts"] } diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index e1c58a0c2b5..da4734eeaf2 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -32,7 +32,6 @@ const config: ViteUserConfig = { "**/*.spec.*", "tests/**", - "e2e_tests/**", "node_modules/**", ".next/**", @@ -44,7 +43,7 @@ const config: ViteUserConfig = { "next.config.*", ], }, - exclude: ["e2e_tests/**", "node_modules/**"], + exclude: ["node_modules/**"], include: ["src/**/*.test.ts", "src/**/*.test.tsx", "tests/**/*.test.ts", "tests/**/*.test.tsx"], }, resolve: { From 38467631b60606693389003b367c785b21637798 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 13:58:15 -0700 Subject: [PATCH 037/130] fix(scim): use members_with_roles as the source of truth for group membership (#34162) * fix(scim): use members_with_roles as the source of truth for group membership SCIM group provisioning tracked membership inconsistently. Team creation and the real team endpoints persist membership in members_with_roles (and each member's user.teams), but the SCIM group PATCH handler and the GET /Groups listing read the legacy team.members String[] column, which team creation never populates. Seeding a PATCH result from that empty column made an Okta "add member" operation recompute the member set from scratch and silently drop everyone already in the team, so users ended up missing from the groups they were provisioned into. Reading the same empty column on GET /Groups reported an empty member list back to the IdP, which drove repeated re-provisioning. Separately, add_new_member appended the team id to user.teams with an unconditional array push. Under the concurrent group PATCHes an IdP sends during a reconcile, each request passed the members_with_roles duplicate check and pushed, so user.teams accumulated duplicate ids for the same team. A duplicate also breaks auth logic that keys off the number of teams a user belongs to. Read current membership from members_with_roles in the SCIM group PATCH seed and the GET /Groups listing, and make the user.teams append idempotent via a filtered update that no-ops once the team is present. Resolves LIT-4283 * fix(scim): address review; atomic user-creation and stop writing legacy members Keep the concurrent-safe team append but create the user via an atomic upsert (create-or-update) instead of a check-then-create, so provisioning the same new user concurrently cannot race into a duplicate-key failure; the team is still appended idempotently by a filtered update so an existing user gets no duplicate team id. Stop writing the legacy team.members column in the group PATCH apply so the only membership record is the source of truth (members_with_roles plus each member's user.teams), reconciled by team_member_add/team_member_delete. Tests: existing add_new_member and team-creation mocks updated to the upsert plus filtered-append shape, and new tests cover atomic creation and that the PATCH apply does not write the legacy members column. --- .../management_endpoints/scim/scim_v2.py | 49 +++--- litellm/proxy/management_helpers/utils.py | 34 ++-- .../scim/test_scim_v2_endpoints.py | 162 ++++++++++++++++++ .../test_team_endpoints.py | 8 + .../test_management_helpers_utils.py | 131 +++++++++++++- 5 files changed, 347 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index fa123b7d76c..9ad221cfffa 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1654,8 +1654,11 @@ async def get_groups( # Convert to SCIM format scim_groups = [] for team in teams: - # Get team members with display names - members = await _get_team_members_display(team.members or []) + # Get team members with display names. members_with_roles is the + # source of truth; the legacy `members` column is not populated by + # team creation, so reading it here would report an empty member + # list to the IdP and trigger repeated re-provisioning. + members = await _get_team_members_display(await _get_team_member_user_ids_from_team(team)) verbose_proxy_logger.debug(f"SCIM GET GROUPS members: {members}") team_alias = getattr(team, "team_alias", team.team_id) team_created_at = team.created_at.isoformat() if team.created_at else None @@ -1885,8 +1888,12 @@ async def _process_group_patch_operations( existing_metadata = existing_team.metadata or {} metadata = dict(existing_metadata) if existing_metadata else {} - # Track member changes - current_members = set(existing_team.members or []) + # Track member changes. members_with_roles is the source of truth for team + # membership; the legacy `members` column is not populated by team creation + # or the real team endpoints, so seeding from it would make an `add`/`remove` + # operation recompute the member set from an empty base and silently drop + # everyone already in the team. + current_members = set(await _get_team_member_user_ids_from_team(existing_team)) final_members = current_members.copy() # Process each patch operation @@ -1963,24 +1970,24 @@ async def _process_group_patch_operations( return update_data, final_members -async def _apply_group_patch_updates( - group_id: str, update_data: Dict[str, Any], final_members: Set[str], prisma_client -): - """Apply patch updates to the group in the database.""" - # Serialize metadata if present +async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any], prisma_client): + """Apply the group's metadata/displayName patch updates to the database. + + Membership itself is not written here; it is reconciled onto the source of + truth (members_with_roles and each member's user.teams) by + _handle_group_membership_changes via team_member_add/team_member_delete. + Writing the legacy `members` column here too would create a second, unread + copy of membership that could drift from the source of truth. + """ if "metadata" in update_data and isinstance(update_data["metadata"], dict): update_data["metadata"] = safe_dumps(update_data["metadata"]) - # Update members list - update_data["members"] = list(final_members) - - # Update team in database - updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": group_id}, - data=update_data, - ) - - return updated_team + if update_data: + return await TeamRepository(prisma_client).table.update( + where={"team_id": group_id}, + data=update_data, + ) + return await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) async def _handle_group_membership_changes(group_id: str, current_members: Set[str], final_members: Set[str]): @@ -2036,8 +2043,8 @@ async def patch_group( # Track current members BEFORE update for comparison current_members = set(await _get_team_member_user_ids_from_team(existing_team)) - # Apply updates to the database - updated_team = await _apply_group_patch_updates(group_id, update_data, final_members, prisma_client) + # Apply the metadata/displayName updates to the database + updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client) # Refresh team data from database to get the latest state after concurrent updates # This prevents race conditions when multiple PATCH requests come in simultaneously diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 11a99caebf5..86beb063667 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -252,6 +252,21 @@ async def _resolve_member_budget_id( return response.budget_id +async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None: + """Append team_id to a user's teams array, only if it is not already present. + + The row-level filter makes the append a no-op once the team is present, so + repeated or concurrent adds of the same team cannot accumulate duplicate + team ids in user.teams (a duplicate also breaks auth logic that keys off the + number of teams a user belongs to). Teams added concurrently for a different + team id are unaffected, since each update filters on its own team id. + """ + await UserRepository(prisma_client).table.update_many( + where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}}, + data={"teams": {"push": [team_id]}}, + ) + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -276,13 +291,16 @@ async def add_new_member( ## ADD TEAM ID, to USER TABLE IF NEW ## if new_member.user_id is not None: new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id) + # Upsert ensures the user row exists atomically (no create race when the + # same new user is provisioned concurrently), seeding teams on create. + # The teams append lives in the filtered update below rather than the + # upsert's update branch so an already-existing user does not get a + # duplicate team id. _returned_user = await UserRepository(prisma_client).table.upsert( where={"user_id": new_member.user_id}, - data={ - "update": {"teams": {"push": [team_id]}}, - "create": {"teams": [team_id], **new_user_defaults}, # type: ignore - }, + data={"create": {"teams": [team_id], **new_user_defaults}, "update": {}}, ) + await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id) if _returned_user is not None: returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif new_member.user_email is not None: @@ -302,12 +320,8 @@ async def add_new_member( returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif len(existing_user_row) == 1: user_info = existing_user_row[0] - _returned_user = await UserRepository(prisma_client).table.update( - where={"user_id": user_info.user_id}, # type: ignore - data={"teams": {"push": [team_id]}}, - ) - if _returned_user is not None: - returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) + await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id) + returned_user = LiteLLM_UserTable(**user_info.model_dump()) elif len(existing_user_row) > 1: raise HTTPException( status_code=400, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 3ff8e2a6886..26a8d1b223b 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -4,14 +4,17 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import ( + LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, + Member, NewUserRequest, NewUserResponse, ProxyException, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( UserProvisionerHelpers, + _apply_group_patch_updates, _extract_group_member_ids, _handle_team_membership_changes, _process_group_patch_operations, @@ -19,6 +22,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( create_group, create_user, delete_group, + get_groups, get_users, get_service_provider_config, patch_group, @@ -2855,3 +2859,161 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): recompute_mock.assert_awaited_once() assert set(recompute_mock.call_args[0][1]) == {"user1"} + + +@pytest.mark.asyncio +async def test_process_group_patch_operations_add_retains_existing_members( + mocker, monkeypatch +): + """A SCIM group ``add`` operation must not drop members already in the team. + + Team membership lives in members_with_roles; team creation leaves the legacy + ``members`` column empty. Seeding the patch result from that empty column + made an ``add`` recompute the member set from scratch and remove everyone + already in the team. The result set must be seeded from members_with_roles so + existing members survive an add of a new one. + """ + + async def mock_get_config(): + return {"litellm_settings": {"scim_upsert_user": True}} + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], # legacy column intentionally empty, as real teams leave it + members_with_roles=[Member(user_id="existing-user", role="user")], + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user"}]) + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + # new-user already exists in the DB + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="new-user") + ) + + _, final_members = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=mock_prisma_client, + ) + + assert final_members == {"existing-user", "new-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_operations_remove_uses_members_with_roles( + mocker, monkeypatch +): + """A ``remove`` op must diff against members_with_roles, so removing one + member leaves the rest of the team intact rather than emptying it.""" + + async def mock_get_config(): + return {"litellm_settings": {"scim_upsert_user": True}} + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="keep-user", role="user"), + Member(user_id="drop-user", role="user"), + ], + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation( + op="remove", path="members", value=[{"value": "drop-user"}] + ) + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="drop-user") + ) + + _, final_members = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=mock_prisma_client, + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_get_groups_reports_members_from_members_with_roles(mocker): + """GET /Groups must report members from members_with_roles (the source of + truth), not the legacy ``members`` column that team creation leaves empty. + Reporting an empty member list makes the IdP repeatedly re-provision.""" + team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], # legacy column empty + members_with_roles=[Member(user_id="member-1", role="user")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team]) + mock_prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + + response = await get_groups(startIndex=1, count=10, filter=None) + + assert [m.value for m in response.Resources[0].members] == ["member-1"] + + +@pytest.mark.asyncio +async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker): + """The group PATCH apply must not write the legacy ``members`` column. + + Membership is reconciled onto the source of truth (members_with_roles and + each member's user.teams) separately; writing the legacy column here too + would create a second, unread copy of membership that can drift from the + source of truth, which is the inconsistency this PR removes. + """ + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + updated = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated) + + result = await _apply_group_patch_updates( + group_id="team-1", + update_data={"team_alias": "Renamed"}, + prisma_client=mock_prisma_client, + ) + + assert result is updated + mock_prisma_client.db.litellm_teamtable.update.assert_awaited_once() + written = mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs["data"] + assert "members" not in written + assert written["team_alias"] == "Renamed" 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 4936191c344..50817b6a4c2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4106,6 +4106,8 @@ async def test_new_team_max_budget_within_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -4247,6 +4249,8 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -4393,6 +4397,8 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table @@ -7245,6 +7251,8 @@ async def test_new_team_soft_budget_validation( } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock team membership table diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index dbbdca65cc6..01e5414a469 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -202,7 +202,8 @@ async def test_add_new_member_clones_default_team_budget_id(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -305,7 +306,8 @@ async def test_add_new_member_budget_duration_only_clones_default_max_budget(): "teams": ["team-dc"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_default_budget_row = MagicMock() @@ -388,7 +390,8 @@ async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -455,7 +458,8 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): "teams": [test_team_id], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) @@ -531,7 +535,8 @@ async def test_add_new_member_persists_budget_duration(): "teams": ["team-dur"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_budget_response = MagicMock() @@ -594,7 +599,8 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): "teams": ["team-dur2"], "user_role": "internal_user", } - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mock_user_response ) mock_budget_response = MagicMock() @@ -997,3 +1003,116 @@ async def test_attach_object_permission_to_dict_with_none_object_permission_id() # Verify no database query was made mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_new_member_appends_team_only_if_absent_for_existing_user(): + """Adding an existing user to a team must append the team id only if it is + not already present. + + add_new_member is the single writer of user.teams for every team add + (/team/member_add, /user/new, SSO, SCIM). An unconditional append let + repeated or concurrent adds accumulate duplicate team ids in user.teams, + which also breaks auth logic that keys off the number of teams a user + belongs to. The append must go through a filtered update that no-ops when + the team is already present, and it must not fall through to creating a new + user row for a user that already exists. + """ + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="existing-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_user_after = MagicMock() + mock_user_after.model_dump.return_value = { + "user_id": "existing-user", + "user_email": None, + "teams": ["team-1"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_after) + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() + # no team default budget and no explicit budget -> no team membership row + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + result_user, _ = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-1", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="admin", + ) + + assert result_user is not None + assert result_user.user_id == "existing-user" + + # the append must be a filtered, idempotent update keyed off the team id, so + # a repeated or concurrent add of a team the user already has is a no-op + mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() + where = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["where"] + assert where["user_id"] == "existing-user" + assert where["NOT"] == {"teams": {"has": "team-1"}} + data = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["data"] + assert data == {"teams": {"push": ["team-1"]}} + + # upsert (not an unconditional teams push) is what ensures the row exists, so + # its update branch must not carry a teams push that would duplicate + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + upsert_update = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["update"] + assert "teams" not in upsert_update + + +@pytest.mark.asyncio +async def test_add_new_member_creates_missing_user_atomically_via_upsert(): + """A brand-new user added to a team must be created via an atomic upsert, not + a separate existence check followed by create. + + Concurrent provisioning of the same new user (which SCIM group reconciles do) + would race a check-then-create into a duplicate-key failure. The upsert seeds + teams on create, and the filtered append is a no-op because the team is + already present on the freshly created row. + """ + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="brand-new-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_created = MagicMock() + mock_created.model_dump.return_value = { + "user_id": "brand-new-user", + "user_email": None, + "teams": ["team-1"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_created) + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() + mock_prisma_client.db.litellm_usertable.create = AsyncMock() + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + result_user, _ = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-1", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="admin", + ) + + assert result_user is not None + assert result_user.user_id == "brand-new-user" + + # existence is established by an atomic upsert (create-or-update), never a + # non-atomic standalone create that could race under concurrent provisioning + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + mock_prisma_client.db.litellm_usertable.create.assert_not_called() + create_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["create"] + assert create_data["teams"] == ["team-1"] From 5abe5f82e18e078fcb535df010c88b25d6736213 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 13:59:02 -0700 Subject: [PATCH 038/130] fix(scim): sync team roster and dedup teams for existing-user email upsert (#34183) * fix(scim): sync team roster and dedup teams for existing-user email upsert When POST /scim/v2/Users matched an already-existing user by email, handle_existing_user_by_email raw-wrote the user's teams array but never touched the team roster, so the user appeared in the group on their profile yet was absent from the team directly (members_with_roles stayed empty). It also did not dedup the teams built from repeated SCIM groups. Route the existing-user team assignment through the same _handle_team_membership_changes / team_member_add path the PUT update_user handler uses, so members_with_roles, LiteLLM_TeamMembership, and the user's teams array stay in sync, and dedup the teams derived from user.groups. The user_id rewrite to the new userName is preserved and sequenced before the roster sync so the roster never references a stale primary key. * fix(scim): surface roster add failures on existing-user email upsert Route the existing-email upsert's roster sync through patch_team_membership with a new opt-in raise_on_error flag so a genuine team_member_add failure propagates instead of being swallowed, and the deduped teams array is only persisted after the roster sync succeeds. Without this, a failed add left the endpoint reporting success while user.teams listed a team members_with_roles never received. The benign already-a-member case stays a no-op even under the strict path, and the flag defaults to False so the PUT update_user, PATCH patch_user, and group callers keep their existing best-effort behavior. SCIM POST is idempotent, so surfacing the error lets the IdP retry and converge. * fix(scim): surface roster removal failures symmetrically with adds Make team_member_delete failures fail loud under the strict roster sync used by the existing-email upsert, mirroring the add path, so a swallowed removal can no longer let the user's teams array drop a team the roster still holds. The idempotent case where the user is already absent from the team stays a no-op, matching how an add treats the user already being in the team. Best-effort behavior is preserved for the default raise_on_error=False callers. --- .../management_endpoints/scim/scim_v2.py | 57 +- .../scim/test_scim_v2_endpoints.py | 703 ++++++++++-------- 2 files changed, 425 insertions(+), 335 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 9ad221cfffa..db2cf2b70dd 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -98,14 +98,27 @@ class UserProvisionerHelpers: if not existing_user: return None - # Update the user + new_teams = list(dict.fromkeys(new_user_request.teams or [])) + + if new_user_request.user_id != existing_user.user_id: + await UserRepository(prisma_client).table.update( + where={"user_id": existing_user.user_id}, + data={"user_id": new_user_request.user_id}, + ) + + await _handle_team_membership_changes( + user_id=new_user_request.user_id, + existing_teams=existing_user.teams or [], + new_teams=new_teams, + raise_on_error=True, + ) + updated_user = await UserRepository(prisma_client).table.update( - where={"user_id": existing_user.user_id}, + where={"user_id": new_user_request.user_id}, data={ - "user_id": new_user_request.user_id, "user_email": new_user_request.user_email, "user_alias": new_user_request.user_alias, - "teams": new_user_request.teams, + "teams": new_teams, "metadata": safe_dumps(new_user_request.metadata), **({"user_role": new_user_request.user_role} if admin_group is not None else {}), }, @@ -440,7 +453,12 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: return members -async def _handle_team_membership_changes(user_id: str, existing_teams: List[str], new_teams: List[str]) -> None: +async def _handle_team_membership_changes( + user_id: str, + existing_teams: List[str], + new_teams: List[str], + raise_on_error: bool = False, +) -> None: """Handle adding/removing user from teams based on changes.""" existing_teams_set = set(existing_teams) new_teams_set = set(new_teams) @@ -453,6 +471,7 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str user_id=user_id, teams_ids_to_add_user_to=list(teams_to_add), teams_ids_to_remove_user_from=list(teams_to_remove), + raise_on_error=raise_on_error, ) @@ -1497,16 +1516,29 @@ def _apply_patch_ops( return update_data, final_team_set +def _is_user_not_in_team_error(exc: HTTPException) -> bool: + """True when team_member_delete reports the user was already absent from the + team, which is the idempotent no-op case for a removal.""" + detail = exc.detail + return isinstance(detail, dict) and detail.get("error") == "User not found in team" + + async def patch_team_membership( user_id: str, teams_ids_to_add_user_to: List[str], teams_ids_to_remove_user_from: List[str], + raise_on_error: bool = False, ) -> bool: """ Add or remove user from teams Handles duplicate membership gracefully (idempotent operation). - If a user is already in a team, that's fine - we don't treat it as an error. + A user already being in a team (on add) or already absent from it (on + remove) is treated as a no-op, not an error. + + When ``raise_on_error`` is True a genuine add or remove failure (anything + other than those idempotent no-ops) propagates instead of being swallowed, + so a caller can avoid persisting a teams array the roster never received. """ for _team_id in teams_ids_to_add_user_to: try: @@ -1521,9 +1553,13 @@ async def patch_team_membership( # Handle duplicate membership gracefully - this is idempotent if e.type == ProxyErrorTypes.team_member_already_in_team: verbose_proxy_logger.debug(f"User {user_id} is already in team {_team_id}, skipping add") + elif raise_on_error: + raise else: verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") except Exception as e: + if raise_on_error: + raise verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") for _team_id in teams_ids_to_remove_user_from: @@ -1532,7 +1568,16 @@ async def patch_team_membership( data=TeamMemberDeleteRequest(team_id=_team_id, user_id=user_id), user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) + except HTTPException as e: + if _is_user_not_in_team_error(e): + verbose_proxy_logger.debug(f"User {user_id} is not in team {_team_id}, skipping remove") + elif raise_on_error: + raise + else: + verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") except Exception as e: + if raise_on_error: + raise verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") return True diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 26a8d1b223b..f458645e51f 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -10,6 +10,7 @@ from litellm.proxy._types import ( Member, NewUserRequest, NewUserResponse, + ProxyErrorTypes, ProxyException, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( @@ -59,9 +60,7 @@ async def test_create_user_existing_user_conflict(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value={"user_id": "existing-user"} - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"}) # Mock the _get_prisma_client_or_raise_exception to return our mock mocker.patch( @@ -235,9 +234,7 @@ async def test_create_user_ingests_entitlements_and_roles(mocker, monkeypatch): }, {"value": "bare-entitlement"}, ] - assert created_metadata["scim_roles"] == [ - {"value": "engineering-admin", "type": "role"} - ] + assert created_metadata["scim_roles"] == [{"value": "engineering-admin", "type": "role"}] @pytest.mark.asyncio @@ -261,9 +258,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey default_params = { "user_role": LitellmUserRoles.PROXY_ADMIN, } - monkeypatch.setattr( - "litellm.default_internal_user_params", default_params, raising=False - ) + monkeypatch.setattr("litellm.default_internal_user_params", default_params, raising=False) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -343,10 +338,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp "BUG: _update_litellm_setting did not update litellm.default_internal_user_params in memory. " "The local variable reassignment (in_memory_var = ...) doesn't propagate back." ) - assert ( - litellm.default_internal_user_params.get("user_role") - == LitellmUserRoles.INTERNAL_USER - ) + assert litellm.default_internal_user_params.get("user_role") == LitellmUserRoles.INTERNAL_USER # Step 3: Create a user via SCIM scim_user = SCIMUser( @@ -442,9 +434,7 @@ async def test_get_users_filters_username_by_exposed_scim_username_for_okta(mock take=10, order={"created_at": "desc"}, ) - mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with( - where=expected_where - ) + mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with(where=expected_where) assert response.totalResults == 1 assert response.Resources[0].id == "internal-user-id" @@ -498,9 +488,7 @@ async def test_get_users_filters_email_value_by_user_email(mocker): take=10, order={"created_at": "desc"}, ) - mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with( - where=expected_where - ) + mock_prisma_client.db.litellm_usertable.count.assert_awaited_once_with(where=expected_where) assert response.totalResults == 1 assert response.Resources[0].id == "internal-user-id" @@ -548,15 +536,12 @@ async def test_handle_existing_user_by_email_no_existing_user(mocker): ) assert result is None - mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( - where={"user_email": "test@example.com"} - ) + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(where={"user_email": "test@example.com"}) @pytest.mark.asyncio async def test_handle_existing_user_by_email_existing_user_updated(mocker): - """Should update existing user and return SCIMUser when user with email exists""" - # Mock existing user - create a proper mock object with attributes + """Should rename the existing user, sync team roster, and return SCIMUser""" existing_user = mocker.MagicMock() existing_user.user_id = "old-user-id" existing_user.user_email = "test@example.com" @@ -564,7 +549,6 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): existing_user.teams = ["old-team"] existing_user.metadata = {"old": "data"} - # Mock updated user updated_user = { "user_id": "new-user-id", "user_email": "test@example.com", @@ -573,7 +557,6 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): "metadata": '{"new": "data"}', } - # Mock SCIM user to be returned mock_scim_user = SCIMUser( schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], id="new-user-id", @@ -585,18 +568,17 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) - # Mock the transformation function mock_transform = mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mock_scim_user), ) + mock_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -611,29 +593,276 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker): prisma_client=mock_prisma_client, new_user_request=new_user_request ) - # Verify the result assert result == mock_scim_user - # Verify database operations - mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( - where={"user_email": "test@example.com"} - ) + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(where={"user_email": "test@example.com"}) - mock_prisma_client.db.litellm_usertable.update.assert_called_once_with( - where={"user_id": "old-user-id"}, - data={ - "user_id": "new-user-id", + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 2 + assert update_calls[0].kwargs == { + "where": {"user_id": "old-user-id"}, + "data": {"user_id": "new-user-id"}, + } + assert update_calls[1].kwargs == { + "where": {"user_id": "new-user-id"}, + "data": { "user_email": "test@example.com", "user_alias": "New Name", "teams": ["new-team"], "metadata": '{"new": "data"}', }, + } + + mock_membership.assert_awaited_once_with( + user_id="new-user-id", + existing_teams=["old-team"], + new_teams=["new-team"], + raise_on_error=True, ) - # Verify transformation was called mock_transform.assert_called_once_with(updated_user) +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocker): + """Existing-email upsert must add the user to the team roster via the shared + team_member_add path and dedup the teams built from repeated SCIM groups. + + Regression: previously the user's ``teams`` array was raw-written (with + duplicates) and the team roster (members_with_roles / LiteLLM_TeamMembership) + was never touched, so the user appeared in the group on their profile but was + absent from the team directly. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "same-id" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + mock_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + + new_user_request = NewUserRequest( + user_id="same-id", + user_email="member@example.com", + user_alias="Member", + teams=["team-a", "team-a", "team-b"], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_membership.assert_awaited_once_with( + user_id="same-id", + existing_teams=[], + new_teams=["team-a", "team-b"], + raise_on_error=True, + ) + + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["where"] == {"user_id": "same-id"} + assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_write(mocker): + """A genuine roster add failure must propagate and must not persist the teams array. + + Regression: the roster sync went through patch_team_membership which swallowed + real team_member_add failures, so the endpoint reported success and wrote a + teams array listing a team the roster never received. The strict path now + surfaces the failure so user.teams and members_with_roles cannot diverge. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_add = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team not found"})), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=["missing-team"], + metadata={}, + auto_create_key=False, + ) + + with pytest.raises(HTTPException): + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_awaited_once() + assert mock_prisma_client.db.litellm_usertable.update.await_count == 0 + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_add_already_member_is_noop(mocker): + """Being already in the team is benign even under the strict path: the upsert + succeeds and the deduped teams array is still persisted.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = [] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_add = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock( + side_effect=ProxyException( + message="already in team", + type=ProxyErrorTypes.team_member_already_in_team.value, + param=None, + code=400, + ) + ), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=["team-x"], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_awaited_once() + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["data"]["teams"] == ["team-x"] + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_write(mocker): + """A genuine roster removal failure must propagate and must not persist the teams array, + symmetrically with add failures, so user.teams cannot drop a team the roster still holds.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["old-team"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_delete = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=HTTPException(status_code=500, detail={"error": "No db connected"})), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + with pytest.raises(HTTPException): + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_delete.assert_awaited_once() + assert mock_prisma_client.db.litellm_usertable.update.await_count == 0 + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noop(mocker): + """A user already absent from the team is the idempotent removal no-op even under the + strict path: the upsert succeeds and the deduped teams array is still persisted.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "uid" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["old-team"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mock_team_member_delete = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=HTTPException(status_code=400, detail={"error": "User not found in team"})), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + + new_user_request = NewUserRequest( + user_id="uid", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_delete.assert_awaited_once() + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["data"]["teams"] == [] + + @pytest.mark.asyncio async def test_handle_team_membership_changes_no_changes(mocker): """Should not call patch_team_membership when existing teams equal new teams""" @@ -766,9 +995,7 @@ async def test_update_user_success(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) # Mock dependencies mocker.patch( @@ -819,11 +1046,7 @@ async def test_update_user_not_found(mocker): ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", - AsyncMock( - side_effect=HTTPException( - status_code=404, detail={"error": "User not found"} - ) - ), + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"})), ) # Should raise ProxyException (which wraps the HTTPException) @@ -868,9 +1091,7 @@ async def test_patch_user_success(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) # Mock dependencies mocker.patch( @@ -907,9 +1128,7 @@ async def test_patch_user_not_found(mocker): """Should raise 404 when user doesn't exist for patch""" patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="replace", path="displayName", value="New Name") - ], + Operations=[SCIMPatchOperation(op="replace", path="displayName", value="New Name")], ) # Mock dependencies to raise HTTPException for user not found @@ -919,11 +1138,7 @@ async def test_patch_user_not_found(mocker): ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", - AsyncMock( - side_effect=HTTPException( - status_code=404, detail={"error": "User not found"} - ) - ), + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"})), ) # Should raise ProxyException (which wraps the HTTPException) @@ -943,9 +1158,7 @@ async def test_get_service_provider_config(mocker): # Verify it returns the correct response assert isinstance(result, SCIMServiceProviderConfig) - assert result.schemas == [ - "urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig" - ] + assert result.schemas == ["urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig"] assert result.patch.supported is True assert result.bulk.supported is False assert result.meta is not None @@ -997,21 +1210,15 @@ async def test_update_group_metadata_serialization_issue(mocker): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock team operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Mock user operations mock_user = mocker.MagicMock() mock_user.user_id = "user1" mock_user.user_email = "user1@example.com" # Add proper string value for user_email mock_user.teams = [group_id] - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mock_user - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock the _get_prisma_client_or_raise_exception to return our mock @@ -1048,9 +1255,7 @@ async def test_update_group_metadata_serialization_issue(mocker): metadata = update_data["metadata"] # The fix should ensure metadata is serialized as a JSON string - assert isinstance( - metadata, str - ), f"metadata should be a JSON string, but got {type(metadata)}" + assert isinstance(metadata, str), f"metadata should be a JSON string, but got {type(metadata)}" # Verify we can parse it back to verify it contains the expected data import json @@ -1107,9 +1312,7 @@ async def test_team_membership_management(mocker): # Check calls for adding members add_calls = [ - call - for call in mock_patch_team_membership.call_args_list - if call[1]["teams_ids_to_add_user_to"] == [group_id] + call for call in mock_patch_team_membership.call_args_list if call[1]["teams_ids_to_add_user_to"] == [group_id] ] assert len(add_calls) == 2 # user3 and user4 @@ -1135,9 +1338,7 @@ async def test_team_membership_management(mocker): # Each call should either add OR remove, not both add_teams = call[1]["teams_ids_to_add_user_to"] remove_teams = call[1]["teams_ids_to_remove_user_from"] - assert (len(add_teams) > 0) != ( - len(remove_teams) > 0 - ) # XOR - one should be empty + assert (len(add_teams) > 0) != (len(remove_teams) > 0) # XOR - one should be empty @pytest.mark.asyncio @@ -1188,9 +1389,7 @@ async def test_update_group_e2e(mocker): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock database operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) # Mock the updated team that gets returned from database updated_team = LiteLLM_TeamTable( @@ -1207,16 +1406,12 @@ async def test_update_group_e2e(mocker): "scim_data": scim_group_update.model_dump(), }, ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=updated_team - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) # Mock user validation (all users exist) mock_user = mocker.MagicMock() mock_user.user_id = "test-user" - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mock_user - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) # Mock dependencies mocker.patch( @@ -1269,29 +1464,19 @@ async def test_update_group_e2e(mocker): assert metadata["scim_data"]["displayName"] == "Updated Team Name" # Verify team membership changes were handled correctly - assert ( - mock_patch_team_membership.call_count == 3 - ) # Remove user1, add user3, add user4 + assert mock_patch_team_membership.call_count == 3 # Remove user1, add user3, add user4 # Check membership changes call_args_list = mock_patch_team_membership.call_args_list # Find remove operation (user1) - remove_calls = [ - call - for call in call_args_list - if call[1]["teams_ids_to_remove_user_from"] == [group_id] - ] + remove_calls = [call for call in call_args_list if call[1]["teams_ids_to_remove_user_from"] == [group_id]] assert len(remove_calls) == 1 assert remove_calls[0][1]["user_id"] == "user1" assert remove_calls[0][1]["teams_ids_to_add_user_to"] == [] # Find add operations (user3, user4) - add_calls = [ - call - for call in call_args_list - if call[1]["teams_ids_to_add_user_to"] == [group_id] - ] + add_calls = [call for call in call_args_list if call[1]["teams_ids_to_add_user_to"] == [group_id]] assert len(add_calls) == 2 add_user_ids = {call[1]["user_id"] for call in add_calls} assert add_user_ids == {"user3", "user4"} @@ -1306,9 +1491,7 @@ async def test_update_group_e2e(mocker): assert len(result.members) == 3 # Verify SCIM transformation was called with updated team - ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with( - updated_team - ) + ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team) @pytest.mark.asyncio @@ -1334,15 +1517,9 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): id=group_id, displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - SCIMMember( - value="new-user-2", display="New User 2" - ), # This user doesn't exist + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist + SCIMMember(value="new-user-2", display="New User 2"), # This user doesn't exist ], ) @@ -1368,9 +1545,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1385,9 +1560,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): # Verify it's a 400 Bad Request assert int(exc_info.value.code) == 400 assert "does not exist" in str(exc_info.value.message) - assert "new-user-1" in str(exc_info.value.message) or "new-user-2" in str( - exc_info.value.message - ) + assert "new-user-1" in str(exc_info.value.message) or "new-user-2" in str(exc_info.value.message) @pytest.mark.asyncio @@ -1422,15 +1595,9 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): id=group_id, displayName="Updated Group Name", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-3", display="New User 3" - ), # This user doesn't exist - SCIMMember( - value="new-user-4", display="New User 4" - ), # This user doesn't exist + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-3", display="New User 3"), # This user doesn't exist + SCIMMember(value="new-user-4", display="New User 4"), # This user doesn't exist ], ) @@ -1441,18 +1608,14 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # Mock team operations - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock updated team response mock_updated_team = mocker.MagicMock() mock_updated_team.team_id = group_id mock_updated_team.team_alias = "Updated Group Name" mock_updated_team.members = ["existing-user", "new-user-3", "new-user-4"] - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team - ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Mock user lookup - only existing-user exists def mock_user_lookup(where): @@ -1463,9 +1626,7 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-3 and new-user-4 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1485,15 +1646,11 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): # Verify it's a 400 Bad Request assert int(exc_info.value.code) == 400 assert "does not exist" in str(exc_info.value.message) - assert "new-user-3" in str(exc_info.value.message) or "new-user-4" in str( - exc_info.value.message - ) + assert "new-user-3" in str(exc_info.value.message) or "new-user-4" in str(exc_info.value.message) @pytest.mark.asyncio -async def test_create_group_with_nonexistent_users_creates_when_flag_true( - mocker, monkeypatch -): +async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker, monkeypatch): """ Test that creating a group with non-existent users creates them when scim_upsert_user is True. This preserves backward compatible behavior. @@ -1514,15 +1671,9 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( id=group_id, displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be created - SCIMMember( - value="new-user-2", display="New User 2" - ), # This user doesn't exist - should be created + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be created + SCIMMember(value="new-user-2", display="New User 2"), # This user doesn't exist - should be created ], ) @@ -1544,9 +1695,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1595,9 +1744,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true( @pytest.mark.asyncio -async def test_extract_group_member_ids_with_flag_true_creates_users( - mocker, monkeypatch -): +async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, monkeypatch): """ Test that _extract_group_member_ids creates users when scim_upsert_user is True. """ @@ -1616,12 +1763,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( id="test-group", displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be created + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be created ], ) @@ -1639,9 +1782,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1666,9 +1807,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users( assert len(result.created_users) == 1 # Verify user was created - mock_create_user.assert_called_once_with( - user_id="new-user-1", created_via="scim_group_membership" - ) + mock_create_user.assert_called_once_with(user_id="new-user-1", created_via="scim_group_membership") @pytest.mark.asyncio @@ -1691,12 +1830,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa id="test-group", displayName="Test Group", members=[ - SCIMMember( - value="existing-user", display="Existing User" - ), # This user exists - SCIMMember( - value="new-user-1", display="New User 1" - ), # This user doesn't exist - should be rejected + SCIMMember(value="existing-user", display="Existing User"), # This user exists + SCIMMember(value="new-user-1", display="New User 1"), # This user doesn't exist - should be rejected ], ) @@ -1714,9 +1849,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=mock_user_lookup - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) # Mock dependencies mocker.patch( @@ -1735,9 +1868,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa @pytest.mark.asyncio -async def test_process_group_patch_operations_with_flag_true_creates_users( - mocker, monkeypatch -): +async def test_process_group_patch_operations_with_flag_true_creates_users(mocker, monkeypatch): """ Test that _process_group_patch_operations creates users when scim_upsert_user is True. """ @@ -1753,11 +1884,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users( # Test data patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="add", path="members", value=[{"value": "new-user-1"}] - ) - ], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user-1"}])], ) # Mock existing team @@ -1791,15 +1918,11 @@ async def test_process_group_patch_operations_with_flag_true_creates_users( assert "new-user-1" in final_members # Verify user was created - mock_create_user.assert_called_once_with( - user_id="new-user-1", created_via="scim_group_patch" - ) + mock_create_user.assert_called_once_with(user_id="new-user-1", created_via="scim_group_patch") @pytest.mark.asyncio -async def test_process_group_patch_operations_with_flag_false_rejects( - mocker, monkeypatch -): +async def test_process_group_patch_operations_with_flag_false_rejects(mocker, monkeypatch): """ Test that _process_group_patch_operations rejects non-existent users when scim_upsert_user is False. """ @@ -1815,11 +1938,7 @@ async def test_process_group_patch_operations_with_flag_false_rejects( # Test data patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="add", path="members", value=[{"value": "new-user-1"}] - ) - ], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user-1"}])], ) # Mock existing team @@ -1894,9 +2013,7 @@ async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeyp @pytest.mark.asyncio -async def test_create_user_keeps_default_when_not_in_scim_admin_group( - mocker, monkeypatch -): +async def test_create_user_keeps_default_when_not_in_scim_admin_group(mocker, monkeypatch): """When scim_admin_group is configured but the user's groups don't include it, the user keeps the non-admin default role.""" from litellm.proxy.proxy_server import proxy_config @@ -1940,9 +2057,7 @@ async def test_create_user_keeps_default_when_not_in_scim_admin_group( @pytest.mark.asyncio -async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( - mocker, monkeypatch -): +async def test_update_user_demotes_admin_when_removed_from_scim_admin_group(mocker, monkeypatch): """Core demotion test: a PUT whose new groups no longer include the configured admin group must re-evaluate the role and write the non-admin default, so an admin removed from the IdP group is demoted without re-login.""" @@ -1976,9 +2091,7 @@ async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2004,9 +2117,7 @@ async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( @pytest.mark.asyncio -async def test_update_user_does_not_force_role_when_scim_admin_group_unset( - mocker, monkeypatch -): +async def test_update_user_does_not_force_role_when_scim_admin_group_unset(mocker, monkeypatch): """When scim_admin_group is unset, PUT must not touch user_role (current behavior preserved).""" from litellm.proxy.proxy_server import proxy_config @@ -2039,9 +2150,7 @@ async def test_update_user_does_not_force_role_when_scim_admin_group_unset( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2067,9 +2176,7 @@ async def test_update_user_does_not_force_role_when_scim_admin_group_unset( @pytest.mark.asyncio -async def test_update_user_demotes_when_default_params_lack_user_role( - mocker, monkeypatch -): +async def test_update_user_demotes_when_default_params_lack_user_role(mocker, monkeypatch): """Regression: default_internal_user_params set without a user_role key must still resolve to the non-admin default on demotion, not silently skip and leave the user PROXY_ADMIN.""" @@ -2079,9 +2186,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} monkeypatch.setattr(proxy_config, "get_config", mock_get_config) - monkeypatch.setattr( - "litellm.default_internal_user_params", {"max_budget": 10}, raising=False - ) + monkeypatch.setattr("litellm.default_internal_user_params", {"max_budget": 10}, raising=False) existing_user = mocker.MagicMock() existing_user.teams = ["litellm-admins"] @@ -2105,9 +2210,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2133,9 +2236,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role( @pytest.mark.asyncio -async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( - mocker, monkeypatch -): +async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group(mocker, monkeypatch): """PATCH that drops the admin team from the resulting team set must write the non-admin default, mirroring the PUT demotion path.""" from litellm.proxy.proxy_server import proxy_config @@ -2152,11 +2253,7 @@ async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="replace", path="groups", value=[{"value": "engineering"}] - ) - ], + Operations=[SCIMPatchOperation(op="replace", path="groups", value=[{"value": "engineering"}])], ) updated_user = { @@ -2172,13 +2269,9 @@ async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=engineering_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=engineering_team) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2227,11 +2320,7 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation( - op="replace", path="groups", value=[{"value": "team-abc-123"}] - ) - ], + Operations=[SCIMPatchOperation(op="replace", path="groups", value=[{"value": "team-abc-123"}])], ) updated_user = { @@ -2247,13 +2336,9 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value=updated_user - ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user) mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=admin_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=admin_team) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2306,9 +2391,7 @@ def _scim_admin_prisma(mocker, *, user_teams): @pytest.mark.asyncio -async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group(mocker, monkeypatch): """The shared recompute helper writes the non-admin default for a member whose resulting teams no longer include the configured admin group.""" from litellm.proxy.proxy_server import proxy_config @@ -2328,9 +2411,7 @@ async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group( @pytest.mark.asyncio -async def test_recompute_scim_member_roles_grants_when_in_admin_group( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_grants_when_in_admin_group(mocker, monkeypatch): """The shared recompute helper grants PROXY_ADMIN when a member's resulting teams include the configured admin group.""" from litellm.proxy.proxy_server import proxy_config @@ -2350,9 +2431,7 @@ async def test_recompute_scim_member_roles_grants_when_in_admin_group( @pytest.mark.asyncio -async def test_recompute_scim_member_roles_noop_when_admin_group_unset( - mocker, monkeypatch -): +async def test_recompute_scim_member_roles_noop_when_admin_group_unset(mocker, monkeypatch): """With scim_admin_group unset the recompute helper must not touch any role, preserving current behavior for SCIM group writes.""" from litellm.proxy.proxy_server import proxy_config @@ -2399,16 +2478,10 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2456,24 +2529,16 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): ) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="remove", path="members", value=[{"value": "user1"}]) - ], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "user1"}])], ) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2523,9 +2588,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) @@ -2556,16 +2619,16 @@ async def test_handle_existing_user_by_email_applies_role_when_admin_group_set(m mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "new-user-id"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "new-user-id"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mocker.MagicMock()), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -2596,16 +2659,16 @@ async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset( mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "new-user-id"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "new-user-id"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=mocker.MagicMock()), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) new_user_request = NewUserRequest( user_id="new-user-id", @@ -2626,9 +2689,7 @@ async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset( @pytest.mark.asyncio -async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( - mocker, monkeypatch -): +async def test_create_user_existing_email_upsert_demotes_when_admin_group_set(mocker, monkeypatch): """End-to-end create wiring: a SCIM POST that upserts an existing email while the user is not in the admin group must write the non-admin default, not leave a stale PROXY_ADMIN.""" @@ -2654,12 +2715,8 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_usertable.update = AsyncMock( - return_value={"user_id": "returning-user"} - ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "returning-user"}) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2673,6 +2730,10 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=scim_user), ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) await create_user(user=scim_user) @@ -2702,9 +2763,7 @@ async def test_create_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2758,16 +2817,10 @@ async def test_update_group_rename_recomputes_retained_members(mocker): mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2812,24 +2865,16 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): ) patch_ops = SCIMPatchOp( schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], - Operations=[ - SCIMPatchOperation(op="replace", path="displayName", value="Engineering") - ], + Operations=[SCIMPatchOperation(op="replace", path="displayName", value="Engineering")], ) mock_prisma_client = mocker.MagicMock() mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=existing_team - ) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock( - return_value=existing_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=mocker.MagicMock() - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", From 482f05190c7ad8b3e478e4d4caff38579367b308 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 14:25:48 -0700 Subject: [PATCH 039/130] fix(scim): prune deleted user from teams' members_with_roles (#34180) SCIM delete_user removed the user from the legacy team.members column and deleted their team membership rows, but never pruned members_with_roles, which is the source of truth ScimTransformations reads for GET /Groups/{id}. A deleted user therefore lingered as a dangling member reference on every team they belonged to Prune each of the user's teams directly via team_member_delete before deleting the user row, and only for teams whose members_with_roles actually contain the user, so a real DB failure surfaces (the endpoint fails loudly and the user is kept; SCIM DELETE is idempotent, so the IdP retries) while a user who was never in a team's members_with_roles stays a no-op. patch_team_membership is left unchanged for its other callers --- .../management_endpoints/scim/scim_v2.py | 7 + .../scim/test_scim_v2_endpoints.py | 124 ++++++++++++++++++ 2 files changed, 131 insertions(+) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index db2cf2b70dd..d1d1c5959d8 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1317,6 +1317,13 @@ async def delete_user( where={"team_id": team.team_id}, data={"members": new_members} ) + team_row = LiteLLM_TeamTable(**team.model_dump()) + if any(member.user_id == user_id for member in team_row.members_with_roles or []): + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=team_row.team_id, user_id=user_id), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + await _set_user_keys_blocked(user_id=user_id, blocked=True) await _delete_rows_referencing_user(prisma_client, user_id=user_id) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index f458645e51f..c31029eb54d 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( create_group, create_user, delete_group, + delete_user, get_groups, get_users, get_service_provider_config, @@ -3062,3 +3063,126 @@ async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker): written = mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs["data"] assert "members" not in written assert written["team_alias"] == "Renamed" + + +def _mock_prisma_for_delete_user(mocker, team): + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.delete = AsyncMock() + return mock_prisma_client + + +def _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user): + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._set_user_keys_blocked", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._delete_rows_referencing_user", + AsyncMock(), + ) + + +@pytest.mark.asyncio +async def test_delete_user_prunes_members_with_roles(mocker): + """Deleting a SCIM user must remove them from every team they belong to via + team_member_delete, which prunes members_with_roles (the source of truth for + SCIM group membership) so GET /Groups no longer returns a dangling reference + to the now-deleted user.""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id, "other-user"], + members_with_roles=[Member(user_id=user_id, role="user"), Member(user_id="other-user", role="admin")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + team_member_delete_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + await delete_user(user_id=user_id) + + team_member_delete_mock.assert_awaited_once() + call = team_member_delete_mock.call_args + assert call.kwargs["data"].team_id == "team-1" + assert call.kwargs["data"].user_id == user_id + assert call.kwargs["user_api_key_dict"].user_role == LitellmUserRoles.PROXY_ADMIN + mock_prisma_client.db.litellm_usertable.delete.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delete_user_surfaces_prune_failure_and_keeps_user(mocker): + """A genuine failure while pruning members_with_roles must surface: the + endpoint fails loudly and the user row is NOT deleted, so we never report a + successful delete while leaving a dangling member (SCIM DELETE is idempotent, + so the IdP retries).""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id], + members_with_roles=[Member(user_id=user_id, role="user")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(side_effect=Exception("database connection lost")), + ) + + with pytest.raises(Exception): + await delete_user(user_id=user_id) + + mock_prisma_client.db.litellm_usertable.delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_user_skips_teams_where_not_a_member(mocker): + """If the user is not in a team's members_with_roles, deletion must treat that + team as a no-op (no team_member_delete call, no error) and still delete the + user, so a stale legacy membership can't block the delete.""" + user_id = "scim-del-user" + + existing_user = mocker.MagicMock() + existing_user.teams = ["team-1"] + + team = LiteLLM_TeamTable( + team_id="team-1", + members=[user_id], + members_with_roles=[Member(user_id="someone-else", role="admin")], + ) + + mock_prisma_client = _mock_prisma_for_delete_user(mocker, team) + _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user) + team_member_delete_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + await delete_user(user_id=user_id) + + team_member_delete_mock.assert_not_awaited() + mock_prisma_client.db.litellm_usertable.delete.assert_awaited_once() From 3f98f6274886ffa511b13e62765f2e2baac96682 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 14:43:13 -0700 Subject: [PATCH 040/130] test(e2e): fail the run when a Rust gateway silently serves /messages through Python (#34208) A gateway whose native extension is unavailable falls back to the Python implementation without raising, so it answers /v1/messages normally and the only difference on the wire is the absent x-litellm-rust header. Nothing in the suite read that header, so a Rust deployment that had stopped running Rust produced a fully green e2e run. Assert the marker on the streamed Messages assertions when E2E_EXPECT_RUST is set. It stays opt-in because the same suite image also runs against the standard gateway, which has no Rust path and must keep passing; the two deployments are already separate Applications, so this is one value on the Rust instance rather than branching inside the tests. --- tests/e2e/e2e_config.py | 2 ++ .../llm_translation/test_messages_azure_foundry_e2e.py | 9 ++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 30353b93dd6..e7c48690c0a 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -72,6 +72,8 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) +EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") + LOAD_USERS = int(os.environ.get("E2E_LOAD_USERS", "750")) LOAD_SPAWN_RATE = float(os.environ.get("E2E_LOAD_SPAWN_RATE", "50")) LOAD_DURATION_SECONDS = float(os.environ.get("E2E_LOAD_DURATION_SECONDS", "60")) diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index 3f2907f202e..d8d44820e80 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -12,7 +12,7 @@ from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import EXPECT_RUST, unique_marker from e2e_http import StreamingResponse, require_successful_call, unwrap from endpoints_client import EndpointsClient from lifecycle import ResourceManager @@ -50,6 +50,13 @@ def _assert_streamed_ok(result: StreamingResponse) -> None: assert any("message_stop" in event for event in result.stream_events), ( "stream never reached message_stop" ) + if EXPECT_RUST: + assert result.headers.get("x-litellm-rust") == "true", ( + "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " + "Rust path, but the response carried no x-litellm-rust marker. The request " + "still succeeded, which is exactly the failure mode: a gateway whose native " + f"extension is unavailable falls back to Python silently. headers={result.headers}" + ) class TestAzureFoundryMessages: From c6b2f111a684b8fc8a7bc2ba27234fcf6b2d1efa Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 14:47:22 -0700 Subject: [PATCH 041/130] fix(team): make team member add atomic to prevent concurrent-add member loss (#34185) _add_team_members_to_team reconciled membership by reading the complete_team_data snapshot captured at the start of team_member_add, appending in memory, and writing the whole members_with_roles array back. Two concurrent /team/member_add calls for the same team read the same snapshot, so the last write wins and one member is silently lost. This affects every concurrent team member add, including the SCIM group PATCH op:add path that routes through team_member_add Reconcile members_with_roles inside a transaction that locks the team row with SELECT ... FOR UPDATE before re-reading the current membership, so concurrent writers serialize on the row lock and each appends onto the other's committed result. The interactive transaction is exposed through a thin PrismaClient.tx() passthrough and the locked read is encapsulated in TeamRepository.get_members_with_roles_locked, and the SCIM group PATCH applies membership as deltas so concurrent adds are not clobbered --- .../management_endpoints/scim/scim_v2.py | 51 +++-- .../management_endpoints/team_endpoints.py | 37 ++-- litellm/proxy/utils.py | 9 + litellm/repositories/team_repository.py | 29 ++- tests/proxy_unit_tests/test_proxy_server.py | 83 +++++--- .../scim/test_scim_v2_endpoints.py | 197 +++++++++++++++++- .../test_team_endpoints.py | 72 +++++++ .../repositories/test_repositories.py | 38 +++- 8 files changed, 448 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index d1d1c5959d8..525ce9f0e89 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1932,8 +1932,16 @@ async def delete_group( async def _process_group_patch_operations( patch_ops: SCIMPatchOp, existing_team, prisma_client -) -> Tuple[Dict[str, Any], Set[str]]: - """Process patch operations for a group and return update data and final members.""" +) -> Tuple[Dict[str, Any], Set[str], Set[str] | None]: + """Process patch operations for a group and return update data, final members + and, when the request contained a member ``replace`` op, the absolute target + roster it declared (``None`` otherwise). + + ``add``/``remove`` are deltas relative to the current roster, but ``replace`` + is absolute: it declares the roster is exactly this set, so the caller must + reconcile against it as a set-to-target rather than rebasing it onto a + concurrently-mutated roster. + """ update_data: Dict[str, Any] = {} # Create a fresh copy of existing metadata to avoid Prisma issues @@ -2019,7 +2027,12 @@ async def _process_group_patch_operations( if metadata: update_data["metadata"] = metadata - return update_data, final_members + member_replace_present = any( + op.op == "replace" and (op.path or "").lower().startswith("members") for op in patch_ops.Operations + ) + replace_target = set(final_members) if member_replace_present else None + + return update_data, final_members, replace_target async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any], prisma_client): @@ -2090,27 +2103,29 @@ async def patch_group( existing_team = await _check_team_exists(group_id) # Process patch operations - update_data, final_members = await _process_group_patch_operations(patch_ops, existing_team, prisma_client) + update_data, final_members, replace_target = await _process_group_patch_operations( + patch_ops, existing_team, prisma_client + ) - # Track current members BEFORE update for comparison - current_members = set(await _get_team_member_user_ids_from_team(existing_team)) + snapshot_members = set(await _get_team_member_user_ids_from_team(existing_team)) + intended_add = final_members - snapshot_members + intended_remove = snapshot_members - final_members # Apply the metadata/displayName updates to the database updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client) - # Refresh team data from database to get the latest state after concurrent updates - # This prevents race conditions when multiple PATCH requests come in simultaneously refreshed_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) - if refreshed_team: - # Re-read current members from refreshed team to account for concurrent updates - refreshed_current_members = set( - await _get_team_member_user_ids_from_team(LiteLLM_TeamTable(**refreshed_team.model_dump())) - ) - # Use the refreshed members for comparison - current_members = refreshed_current_members + refreshed_current = ( + set(await _get_team_member_user_ids_from_team(LiteLLM_TeamTable(**refreshed_team.model_dump()))) + if refreshed_team + else snapshot_members + ) - # Handle user-team relationship changes - await _handle_group_membership_changes(group_id, current_members, final_members) + effective_final = ( + replace_target if replace_target is not None else (refreshed_current | intended_add) - intended_remove + ) + + await _handle_group_membership_changes(group_id, refreshed_current, effective_final) # A rename can flip whether this group matches scim_admin_group by display # name, so retained members must be re-resolved too, not just the ones whose @@ -2119,7 +2134,7 @@ async def patch_group( alias_changed = new_alias != existing_team.team_alias await _recompute_scim_member_roles( prisma_client, - (current_members | final_members if alias_changed else current_members ^ final_members), + (refreshed_current | effective_final if alias_changed else refreshed_current ^ effective_final), ) # Refresh team one more time to get final state after membership changes diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 70c002d2d2d..3a7e89aa20d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2375,7 +2375,15 @@ async def _add_team_members_to_team( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> Tuple[LiteLLM_TeamTable, List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: - """Add team members to the team.""" + """Add team members to the team. + + The members_with_roles reconciliation runs inside a transaction that locks + the team row with ``SELECT ... FOR UPDATE`` before reading the current + membership. Concurrent /team/member_add calls for the same team therefore + serialize on the row lock and each appends onto the other's committed + result, instead of both rewriting the whole JSON array from a stale + snapshot (which silently drops one member on the losing write). + """ # Process and add new members updated_users, updated_team_memberships = await _process_team_members( data=data, @@ -2385,19 +2393,22 @@ async def _add_team_members_to_team( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - # Update team members list - await _update_team_members_list( - data=data, - complete_team_data=complete_team_data, - updated_users=updated_users, - ) + async with prisma_client.tx() as tx: + complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked( + tx, data.team_id + ) - # ADD MEMBER TO TEAM - _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore - ) + await _update_team_members_list( + data=data, + complete_team_data=complete_team_data, + updated_users=updated_users, + ) + + _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] + updated_team = await tx.litellm_teamtable.update( + where={"team_id": data.team_id}, + data={"members_with_roles": json.dumps(_db_team_members)}, + ) return updated_team, updated_users, updated_team_memberships diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 43921a847a9..1dbc0ad1837 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -172,6 +172,7 @@ from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -2922,6 +2923,14 @@ class PrismaClient: return self.db.writer return self.db + def tx(self) -> "TransactionManager": + """Open an interactive transaction on the writer. + + Callers go through this instead of reaching into ``self.db`` so writer + selection and read-replica routing stay encapsulated in the wrapper. + """ + return cast("TransactionManager", self.db.tx()) # cast-ok: wrappers delegate tx via __getattr__ (untyped) + def get_request_status(self, payload: Union[dict, SpendLogsPayload]) -> Literal["success", "failure"]: """ Determine if a request was successful or failed based on payload metadata. diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 3227aa812ca..68875bd7972 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -4,11 +4,18 @@ Team repository for database operations on LiteLLM_TeamTable. import json from datetime import datetime -from typing import Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type -from litellm.models.team import LiteLLM_TeamTable +from pydantic import TypeAdapter + +from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import BaseRepository +if TYPE_CHECKING: + from prisma import Prisma + +_MEMBERS_WITH_ROLES_ADAPTER = TypeAdapter(list[Member]) + class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" @@ -46,6 +53,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable(**data) + async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> List[Member]: + """Return the team's members_with_roles, locking the row FOR UPDATE. + + Must be called inside a transaction so the row lock is held until + commit. This serializes concurrent membership writers on the team row + so the losing writer appends onto the winner's committed result instead + of overwriting it from a stale snapshot. + """ + rows = await tx.query_raw( + 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', + team_id, + ) + raw_value = rows[0]["members_with_roles"] if rows else None + parsed = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + if not parsed: + return [] + return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed) + async def find_by_id(self, team_id: str, id_field: str = "team_id") -> Optional[LiteLLM_TeamTable]: return await super().find_by_id(team_id, id_field) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 212f7772cad..bedd4dd1838 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1252,6 +1252,17 @@ async def test_create_team_member_add(prisma_client, new_member_method): return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) + tx_mock = AsyncMock() + tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + tx_mock.litellm_teamtable = team_mock_client + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) + tx_cm.__aexit__ = AsyncMock(return_value=None) + original_tx = litellm.proxy.proxy_server.prisma_client.tx + litellm.proxy.proxy_server.prisma_client.tx = MagicMock( + return_value=tx_cm + ) + print(f"team_member_add_request={team_member_add_request}") await team_member_add( data=team_member_add_request, @@ -1273,6 +1284,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val + litellm.proxy.proxy_server.prisma_client.tx = original_tx @pytest.mark.parametrize("team_member_role", ["admin", "user"]) @@ -1434,42 +1446,51 @@ async def test_create_team_member_add_team_admin( mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() - original_val = getattr( - litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable" - ) - litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = team_mock_client - team_mock_client.update = AsyncMock( return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) - try: - await team_member_add( - data=team_member_add_request, - user_api_key_dict=valid_token, + tx_mock = AsyncMock() + tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + tx_mock.litellm_teamtable = team_mock_client + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + with ( + patch.object( + litellm.proxy.proxy_server.prisma_client.db, + "litellm_teamtable", + team_mock_client, + ), + patch.object( + litellm.proxy.proxy_server.prisma_client, + "tx", + MagicMock(return_value=tx_cm), + ), + ): + try: + await team_member_add( + data=team_member_add_request, + user_api_key_dict=valid_token, + ) + except HTTPException as e: + if user_role == "user": + assert e.status_code == 403 + return + else: + raise e + + mock_client.assert_called() + + assert ( + mock_client.call_args.kwargs["data"]["create"]["max_budget"] + == litellm.max_internal_user_budget + ) + assert ( + mock_client.call_args.kwargs["data"]["create"]["budget_duration"] + == litellm.internal_user_budget_duration ) - except HTTPException as e: - if user_role == "user": - assert e.status_code == 403 - return - else: - raise e - - mock_client.assert_called() - - print(f"mock_client.call_args: {mock_client.call_args}") - print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs)) - - assert ( - mock_client.call_args.kwargs["data"]["create"]["max_budget"] - == litellm.max_internal_user_budget - ) - assert ( - mock_client.call_args.kwargs["data"]["create"]["budget_duration"] - == litellm.internal_user_budget_duration - ) - - litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index c31029eb54d..850b8fc7cfd 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1909,7 +1909,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users(mocke ) # Execute the function - update_data, final_members = await _process_group_patch_operations( + update_data, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, existing_team=mock_existing_team, prisma_client=mock_prisma_client, @@ -2948,7 +2948,7 @@ async def test_process_group_patch_operations_add_retains_existing_members( return_value=mocker.MagicMock(user_id="new-user") ) - _, final_members = await _process_group_patch_operations( + _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, existing_team=existing_team, prisma_client=mock_prisma_client, @@ -2996,7 +2996,7 @@ async def test_process_group_patch_operations_remove_uses_members_with_roles( return_value=mocker.MagicMock(user_id="drop-user") ) - _, final_members = await _process_group_patch_operations( + _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, existing_team=existing_team, prisma_client=mock_prisma_client, @@ -3186,3 +3186,194 @@ async def test_delete_user_skips_teams_where_not_a_member(mocker): team_member_delete_mock.assert_not_awaited() mock_prisma_client.db.litellm_usertable.delete.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): + """A group PATCH op:add must be applied as a delta against the live roster, + not as a snapshot-based absolute target. + + When a concurrent PATCH has already added a member between this request's + initial read and its post-write refresh, that member shows up in the + refreshed roster but not in this request's snapshot-derived target. Diffing + the refreshed roster against the snapshot target would issue a spurious + team_member_delete for the concurrently-added member. Applying only this + request's intended delta on top of the refreshed roster must retain them. + """ + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "team-concurrent" + + snapshot_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="zed", role="user")], + metadata={"externalId": "grp-ext"}, + ) + refreshed_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="zed", role="user"), + Member(user_id="alice", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + final_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="zed", role="user"), + Member(user_id="alice", role="user"), + Member(user_id="bob", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="add", path="members", value=[{"value": "bob"}])], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=[snapshot_team, refreshed_team, final_team] + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + patch_membership_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Group", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + calls = patch_membership_mock.call_args_list + + removed_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_remove_user_from") == [group_id] + } + assert removed_user_ids == set() + + added_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_add_user_to") == [group_id] + } + assert added_user_ids == {"bob"} + + +@pytest.mark.asyncio +async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mocker): + """A group PATCH ``replace`` op declares the roster is exactly the given set, + so it must reconcile as a set-to-target, not as a delta. + + Unlike ``add``/``remove``, ``replace`` is absolute. A member that another + request added concurrently is present in the refreshed roster but not in the + replace target, and ``replace`` must drop it. Rebasing the replace onto the + refreshed roster (the delta behavior correct only for add/remove) would + wrongly retain that concurrently-added member. + """ + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "team-replace-concurrent" + + snapshot_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="zed", role="user")], + metadata={"externalId": "grp-ext"}, + ) + refreshed_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[ + Member(user_id="alice", role="user"), + Member(user_id="bob", role="user"), + ], + metadata={"externalId": "grp-ext"}, + ) + final_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Group", + members_with_roles=[Member(user_id="alice", role="user")], + metadata={"externalId": "grp-ext"}, + ) + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="replace", path="members", value=[{"value": "alice"}])], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=[snapshot_team, refreshed_team, final_team] + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + patch_membership_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Group", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + calls = patch_membership_mock.call_args_list + + removed_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_remove_user_from") == [group_id] + } + assert removed_user_ids == {"bob"} + + added_user_ids = { + call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_add_user_to") == [group_id] + } + assert added_user_ids == set() 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 50817b6a4c2..d5d61341c93 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1693,6 +1693,78 @@ async def test_update_team_members_list_duplicate_prevention(): assert len(mock_team.members_with_roles) == 1 +@pytest.mark.asyncio +async def test_add_team_members_reconciles_against_freshly_locked_row(): + """ + Regression: _add_team_members_to_team must build the new members_with_roles + from the row it re-reads under a lock inside the write transaction, not from + the stale complete_team_data snapshot captured at the start of the request. + + Two concurrent /team/member_add calls for the same team read the same + snapshot; without the locked re-read the losing write rewrites the whole + JSON array from its stale copy and silently drops the member the other call + already committed. Here the snapshot holds only "zed", a concurrent writer + has already committed "alice" (returned by the locked SELECT), and this call + adds "bob". The write must contain all three. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + stale_snapshot = LiteLLM_TeamTable( + team_id="test-team-lock", + members_with_roles=[Member(user_id="zed", role="user")], + ) + + freshly_committed = [ + {"user_id": "zed", "user_email": None, "role": "user"}, + {"user_id": "alice", "user_email": None, "role": "user"}, + ] + + captured: dict = {} + + async def _capture_update(where, data): + captured["data"] = data + return LiteLLM_TeamTable( + team_id="test-team-lock", + members_with_roles=json.loads(data["members_with_roles"]), + ) + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": freshly_committed}]) + tx.litellm_teamtable.update = AsyncMock(side_effect=_capture_update) + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([], [])), + ): + updated_team, _, _ = await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="test-team-lock", + member=Member(user_id="bob", role="user"), + ), + complete_team_data=stale_snapshot, + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + written_ids = sorted(m["user_id"] for m in json.loads(captured["data"]["members_with_roles"])) + assert written_ids == ["alice", "bob", "zed"] + + lock_reads = [call for call in tx.query_raw.call_args_list if "FOR UPDATE" in str(call.args[0])] + assert lock_reads, "expected a SELECT ... FOR UPDATE row-lock read before the write" + + assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] + + def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index af2eea823f4..6308faf8fc7 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -5,7 +5,7 @@ Tests for gateway repository layer. import json from datetime import datetime from typing import Any, Dict, List, Optional -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -499,6 +499,42 @@ class TestTeamRepository: assert team.team_id == "team-123" assert team.team_alias == "Engineering" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raw_value, expected_ids", + [ + ( + [ + {"user_id": "a", "role": "user"}, + {"user_id": "b", "role": "admin"}, + ], + ["a", "b"], + ), + (json.dumps([{"user_id": "a", "role": "user"}]), ["a"]), + ({}, []), + (None, []), + ], + ) + async def test_get_members_with_roles_locked(self, repo, raw_value, expected_ids): + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": raw_value}]) + + members = await repo.get_members_with_roles_locked(tx, "team-1") + + assert [m.user_id for m in members] == expected_ids + sql = tx.query_raw.call_args.args[0] + assert "FOR UPDATE" in sql + assert tx.query_raw.call_args.args[1] == "team-1" + + @pytest.mark.asyncio + async def test_get_members_with_roles_locked_missing_row(self, repo): + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[]) + + members = await repo.get_members_with_roles_locked(tx, "missing") + + assert members == [] + @pytest.mark.asyncio async def test_create_team_all_fields(self, repo): team = await repo.create_team( From 9baea68f37d60e2b7c3be33841490600a5c13d4d Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 22 Jul 2026 14:47:35 -0700 Subject: [PATCH 042/130] fix(ui): resolve SSO and SMTP settings from a typed config object (#33576) The SSO and Email Server settings pages read only stored config, so a gateway configured entirely through environment variables rendered every field blank even though both features were live. Rather than add per-endpoint env fallback, resolve each setting through one typed config object. A FieldDescriptor names, for one setting, where it lives in the stored row (db_key), which process env var carries it (env_var), whether it is a secret, and its effective default. A pure resolve_fields reconciles a descriptor table against the stored row and the process environment with a fixed precedence and reports per-field provenance (db, env, default, or unset). The SSO descriptor table single-sources the field-to-env mapping that the read and write paths previously duplicated, so they can no longer drift. get_sso_settings and the /get/config/callbacks alerting block read through the resolver instead of their own inline fallbacks. get_sso_settings no longer decrypts stored values into os.environ; decryption happens once inside the resolver via the pure helper, so a GET stops mutating the process environment. The SSO response carries provenance so the UI can distinguish an env-sourced value from a stored one, and secrets are masked at the endpoint (the resolver returns them unmasked so the login path could consume them). os.environ remains the runtime carrier; the SSO login and mail-send paths are unchanged. The settings pages also submit only fields an admin actually edited, so a rendered mask or env-sourced value is never written back over a working secret, and generic_scope is a real SSO form field. Omitting a field from /update/sso_settings clears it, which provider switching relies on; the deeper write-path concern that behaviour points at is tracked in LIT-4498. --- .../workflows/test-unit-proxy-endpoints.yml | 1 + litellm/proxy/config_resolvers/__init__.py | 9 + .../proxy/config_resolvers/_descriptors.py | 73 ++++++++ litellm/proxy/config_resolvers/alerting.py | 25 +++ litellm/proxy/config_resolvers/sso.py | 94 ++++++++++ litellm/proxy/proxy_server.py | 38 ++--- .../proxy_setting_endpoints.py | 101 ++--------- .../proxy/management_endpoints/ui_sso.py | 4 + .../config_resolvers/test_config_resolvers.py | 105 ++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 107 ++++++++++++ .../test_proxy_setting_endpoints.py | 161 ++++++++++++++++-- .../(dashboard)/hooks/sso/useSSOSettings.ts | 1 + .../src/components/SSOModals.test.tsx | 1 + .../src/components/SSOModals.tsx | 97 +---------- .../Modals/BaseSSOSettingsForm.test.tsx | 63 ++++++- .../Modals/BaseSSOSettingsForm.tsx | 7 +- .../AdminSettings/SSOSettings/SSOSettings.tsx | 2 + .../src/components/email_settings.tsx | 12 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 + 19 files changed, 693 insertions(+), 217 deletions(-) create mode 100644 litellm/proxy/config_resolvers/__init__.py create mode 100644 litellm/proxy/config_resolvers/_descriptors.py create mode 100644 litellm/proxy/config_resolvers/alerting.py create mode 100644 litellm/proxy/config_resolvers/sso.py create mode 100644 tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index cbb36eebdb9..b3eb8f79a43 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -46,6 +46,7 @@ jobs: tests/test_litellm/proxy/rag_endpoints tests/test_litellm/proxy/realtime_endpoints tests/test_litellm/proxy/ui_crud_endpoints + tests/test_litellm/proxy/config_resolvers tests/test_litellm/proxy/utils workers: 2 reruns: 2 diff --git a/litellm/proxy/config_resolvers/__init__.py b/litellm/proxy/config_resolvers/__init__.py new file mode 100644 index 00000000000..88b4c3961f0 --- /dev/null +++ b/litellm/proxy/config_resolvers/__init__.py @@ -0,0 +1,9 @@ +"""Typed, provenance-aware resolution of proxy settings from DB then env.""" + +from litellm.proxy.config_resolvers._descriptors import ( + FieldDescriptor, + FieldSource, + resolve_fields, +) + +__all__ = ["FieldDescriptor", "FieldSource", "resolve_fields"] diff --git a/litellm/proxy/config_resolvers/_descriptors.py b/litellm/proxy/config_resolvers/_descriptors.py new file mode 100644 index 00000000000..f67a690f92f --- /dev/null +++ b/litellm/proxy/config_resolvers/_descriptors.py @@ -0,0 +1,73 @@ +"""Shared primitive for resolving a settings value from its sources. + +A ``FieldDescriptor`` names, for one setting, where it lives in the stored DB +row (``db_key``), which process env var carries it (``env_var``), whether it is +a secret, and its effective default. ``resolve_fields`` reconciles a set of +descriptors against a decrypted DB row and the process environment with a fixed +precedence, returning the resolved values plus per-field provenance so a caller +can tell whether a value came from the database, the environment, a default, or +is unset. +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal + +FieldSource = Literal["db", "env", "default", "unset"] + + +@dataclass(frozen=True, slots=True) +class FieldDescriptor: + field_name: str + db_key: str + env_var: str + is_secret: bool = False + default: str | None = None + + +def _db_is_set(db_value: object, empty_db_is_set: bool) -> bool: + if empty_db_is_set: + # A stored key that is present, even as "", is an explicit admin choice + # (e.g. clearing an alerting webhook) and must win over a stale env var. + return db_value is not None + # A blank stored value is treated as absent, so it falls through to env. This + # fits settings whose clear path also unsets the env var (e.g. SSO). + return isinstance(db_value, str) and bool(db_value.strip()) + + +def _resolve_one( + descriptor: FieldDescriptor, + db_values: Mapping[str, object], + env: Mapping[str, str], + empty_db_is_set: bool, +) -> tuple[str, str | None, FieldSource]: + db_value = db_values.get(descriptor.db_key) + if _db_is_set(db_value, empty_db_is_set): + return descriptor.field_name, db_value if isinstance(db_value, str) else str(db_value), "db" + env_value = env.get(descriptor.env_var) + if isinstance(env_value, str) and env_value.strip(): + return descriptor.field_name, env_value, "env" + if descriptor.default is not None: + return descriptor.field_name, descriptor.default, "default" + return descriptor.field_name, None, "unset" + + +def resolve_fields( + descriptors: Sequence[FieldDescriptor], + db_values: Mapping[str, object], + env: Mapping[str, str], + empty_db_is_set: bool = False, +) -> tuple[dict[str, str | None], dict[str, FieldSource]]: + """Resolve every descriptor to (values, provenance). + + Precedence per field: a set stored value wins, else a non-blank process env + var, else the descriptor default, else unset. ``empty_db_is_set`` selects + how a present-but-empty stored value is read: ``False`` treats it as absent + so it falls back to env (SSO, whose clear path also unsets the env var); + ``True`` treats it as an explicit clear that wins over env (alerting, whose + clear path stores "" without unsetting the env var). + """ + resolved = tuple(_resolve_one(descriptor, db_values, env, empty_db_is_set) for descriptor in descriptors) + values = {field_name: value for field_name, value, _ in resolved} + provenance = {field_name: source for field_name, _, source in resolved} + return values, provenance diff --git a/litellm/proxy/config_resolvers/alerting.py b/litellm/proxy/config_resolvers/alerting.py new file mode 100644 index 00000000000..3704ec09355 --- /dev/null +++ b/litellm/proxy/config_resolvers/alerting.py @@ -0,0 +1,25 @@ +"""Descriptor tables for the alerting settings surfaced by /get/config/callbacks. + +These reconcile the stored ``environment_variables`` blob (keyed by the +uppercase env-var names) with the process environment. SMTP_PORT and SMTP_TLS +carry the same effective defaults the mail-send path applies, so the settings +page shows the config that mail would actually use rather than a blank. +""" + +from litellm.proxy.config_resolvers._descriptors import FieldDescriptor + +EMAIL_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("SMTP_HOST", "SMTP_HOST", "SMTP_HOST"), + FieldDescriptor("SMTP_PORT", "SMTP_PORT", "SMTP_PORT", default="587"), + FieldDescriptor("SMTP_TLS", "SMTP_TLS", "SMTP_TLS", default="True"), + FieldDescriptor("SMTP_USERNAME", "SMTP_USERNAME", "SMTP_USERNAME", is_secret=True), + FieldDescriptor("SMTP_PASSWORD", "SMTP_PASSWORD", "SMTP_PASSWORD", is_secret=True), + FieldDescriptor("SMTP_SENDER_EMAIL", "SMTP_SENDER_EMAIL", "SMTP_SENDER_EMAIL"), + FieldDescriptor("TEST_EMAIL_ADDRESS", "TEST_EMAIL_ADDRESS", "TEST_EMAIL_ADDRESS"), + FieldDescriptor("EMAIL_LOGO_URL", "EMAIL_LOGO_URL", "EMAIL_LOGO_URL"), + FieldDescriptor("EMAIL_SUPPORT_CONTACT", "EMAIL_SUPPORT_CONTACT", "EMAIL_SUPPORT_CONTACT"), +) + +SLACK_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("SLACK_WEBHOOK_URL", "SLACK_WEBHOOK_URL", "SLACK_WEBHOOK_URL", is_secret=True), +) diff --git a/litellm/proxy/config_resolvers/sso.py b/litellm/proxy/config_resolvers/sso.py new file mode 100644 index 00000000000..3d83c06dd62 --- /dev/null +++ b/litellm/proxy/config_resolvers/sso.py @@ -0,0 +1,94 @@ +"""Resolved SSO config object. + +Reconciles the dedicated ``sso_config`` DB row (lowercase, per-value encrypted +keys) with the process environment (uppercase env vars) into a typed +``SSOConfig`` plus per-field provenance. This is the single source of truth for +the SSO field -> env-var mapping, used by both the read-back endpoint and the +save endpoint so the two can never drift. +""" + +from collections.abc import Mapping +from dataclasses import dataclass + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.config_resolvers._descriptors import ( + FieldDescriptor, + FieldSource, + resolve_fields, +) +from litellm.types.proxy.management_endpoints.ui_sso import ( + RoleMappings, + SSOConfig, + TeamMappings, +) + +SSO_DESCRIPTORS: tuple[FieldDescriptor, ...] = ( + FieldDescriptor("google_client_id", "google_client_id", "GOOGLE_CLIENT_ID"), + FieldDescriptor("google_client_secret", "google_client_secret", "GOOGLE_CLIENT_SECRET", is_secret=True), + FieldDescriptor("microsoft_client_id", "microsoft_client_id", "MICROSOFT_CLIENT_ID"), + FieldDescriptor("microsoft_client_secret", "microsoft_client_secret", "MICROSOFT_CLIENT_SECRET", is_secret=True), + FieldDescriptor("microsoft_tenant", "microsoft_tenant", "MICROSOFT_TENANT"), + FieldDescriptor("generic_client_id", "generic_client_id", "GENERIC_CLIENT_ID"), + FieldDescriptor("generic_client_secret", "generic_client_secret", "GENERIC_CLIENT_SECRET", is_secret=True), + FieldDescriptor( + "generic_authorization_endpoint", "generic_authorization_endpoint", "GENERIC_AUTHORIZATION_ENDPOINT" + ), + FieldDescriptor("generic_token_endpoint", "generic_token_endpoint", "GENERIC_TOKEN_ENDPOINT"), + FieldDescriptor("generic_userinfo_endpoint", "generic_userinfo_endpoint", "GENERIC_USERINFO_ENDPOINT"), + FieldDescriptor("generic_scope", "generic_scope", "GENERIC_SCOPE", default="openid email profile"), + FieldDescriptor("proxy_base_url", "proxy_base_url", "PROXY_BASE_URL"), +) + +# Derived from the descriptor table so read (masking) and the field->env mapping +# never diverge from the resolver. +SSO_SECRET_FIELDS: frozenset[str] = frozenset(d.field_name for d in SSO_DESCRIPTORS if d.is_secret) +SSO_FIELD_ENV_VARS: dict[str, str] = {d.field_name: d.env_var for d in SSO_DESCRIPTORS} + +# Structured sub-objects stored on the SSO row that are not simple env-backed +# scalars; handled outside the descriptor resolution. +_STRUCTURED_KEYS = ("role_mappings", "team_mappings") + + +@dataclass(frozen=True, slots=True) +class ResolvedSSOConfig: + config: SSOConfig + provenance: dict[str, FieldSource] + + +def _decrypt(raw: Mapping[str, object]) -> dict[str, object]: + return { + key: ( + decrypt_value_helper(value=value, key=key, return_original_value=True) if isinstance(value, str) else value + ) + for key, value in raw.items() + } + + +def _parse_role_mappings(data: object) -> RoleMappings | None: + # The stored row is JSON, so mappings arrive as a dict (or are absent). + return RoleMappings(**data) if isinstance(data, dict) else None + + +def _parse_team_mappings(data: object) -> TeamMappings | None: + return TeamMappings(**data) if isinstance(data, dict) else None + + +def resolve_sso_config(sso_db_settings: Mapping[str, object] | None, env: Mapping[str, str]) -> ResolvedSSOConfig: + """Resolve the effective SSO config: stored row first, then process env. + + Decryption happens here, once, via the pure ``decrypt_value_helper``; this + function never writes ``os.environ`` (unlike the legacy read path). Values + are returned unmasked so the login path could consume them; the read-back + endpoint is responsible for masking secrets before responding to the UI. + """ + raw = dict(sso_db_settings) if sso_db_settings else {} + decrypted = _decrypt({key: value for key, value in raw.items() if key not in _STRUCTURED_KEYS}) + values, provenance = resolve_fields(SSO_DESCRIPTORS, decrypted, env) + structured = { + "user_email": decrypted.get("user_email"), + "ui_access_mode": decrypted.get("ui_access_mode"), + "role_mappings": _parse_role_mappings(raw.get("role_mappings")), + "team_mappings": _parse_team_mappings(raw.get("team_mappings")), + } + config = SSOConfig(**{**values, **structured}) + return ResolvedSSOConfig(config=config, provenance=provenance) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7677ae51f9a..32845763f22 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -304,6 +304,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form +from litellm.proxy.config_resolvers import resolve_fields +from litellm.proxy.config_resolvers.alerting import ( + EMAIL_DESCRIPTORS, + SLACK_DESCRIPTORS, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -1159,9 +1164,9 @@ _OPENAPI_HTTP_METHODS = { # Credentials surfaced by `/get/config/callbacks` in the alerting block: the # full Slack incoming-webhook URL is itself a credential, and the SMTP # password is a service password. Masked on read so plaintext never reaches -# the UI. Kept here at module scope to match the analogous -# `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO -# and cache endpoint files. +# the UI. Kept here at module scope to match the analogous descriptor +# `is_secret` flags in litellm.proxy.config_resolvers and the +# `_CACHE_SENSITIVE_FIELDS` constant in the cache endpoint file. _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} @@ -15491,14 +15496,10 @@ async def get_config( _alerting = _general_settings.get("alerting", []) alerting_data = [] if "slack" in _alerting: - _slack_vars = [ - "SLACK_WEBHOOK_URL", - ] - _slack_env_vars = { - _var: (value if (value := environment_variables.get(_var)) is not None else os.getenv(_var)) - for _var in _slack_vars - } - _slack_env_vars = _apply_alerting_env_role_gate(_slack_env_vars, is_full_admin) + _slack_values, _ = resolve_fields( + SLACK_DESCRIPTORS, environment_variables, os.environ, empty_db_is_set=True + ) + _slack_env_vars = _apply_alerting_env_role_gate(_slack_values, is_full_admin) _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types _all_alert_types = proxy_logging_obj.slack_alerting_instance._all_possible_alert_types() @@ -15514,19 +15515,8 @@ async def get_config( } ) # pass email alerting vars - _email_vars = [ - "SMTP_HOST", - "SMTP_PORT", - "SMTP_USERNAME", - "SMTP_PASSWORD", - "SMTP_SENDER_EMAIL", - "TEST_EMAIL_ADDRESS", - "EMAIL_LOGO_URL", - "EMAIL_SUPPORT_CONTACT", - ] - _email_env_vars = _apply_alerting_env_role_gate( - {_var: environment_variables.get(_var) for _var in _email_vars}, is_full_admin - ) + _email_values, _ = resolve_fields(EMAIL_DESCRIPTORS, environment_variables, os.environ, empty_db_is_set=True) + _email_env_vars = _apply_alerting_env_role_gate(_email_values, is_full_admin) alerting_data.append( { diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 42111cf17f2..10c71c00110 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -15,6 +15,11 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers.sso import ( + SSO_FIELD_ENV_VARS, + SSO_SECRET_FIELDS, + resolve_sso_config, +) from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( SSOConfigRepository, @@ -27,16 +32,6 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router = APIRouter() -# SSO secret fields returned by /get/sso_settings. These are masked on read so -# the UI can show "(set)" without ever transporting the plaintext OAuth secret -# off the server, matching the write-once + masked-on-read contract used for -# the HashiCorp Vault config override. -_SSO_SENSITIVE_FIELDS: Set[str] = { - "google_client_secret", - "microsoft_client_secret", - "generic_client_secret", -} - # Maps each UIThemeConfig field to the env var the UI branding path reads it # from. /update/ui_theme_settings writes both the stored ui_theme_config and # these env vars, so /get/ui_theme_settings resolves the same env vars to @@ -109,7 +104,8 @@ class SettingsResponse(BaseModel): class SSOSettingsResponse(SettingsResponse): """Response model for SSO settings""" - pass + provenance: Dict[str, str] = Field(default_factory=dict) + """Per-field source of each value: 'db', 'env', 'default', or 'unset'.""" class InternalUserSettingsResponse(SettingsResponse): @@ -757,7 +753,7 @@ async def get_sso_settings(): Returns a structured object with values and descriptions for UI display. """ - from litellm.proxy.proxy_server import prisma_client, proxy_config + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException( @@ -765,59 +761,12 @@ async def get_sso_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - # Get SSO config from dedicated table + # Resolve the effective SSO config: the stored row wins, else the process + # environment, else each field's default. Unlike the legacy read path this + # does not write os.environ; a GET has no business mutating the environment. sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - - # Initialize with defaults - sso_settings_dict = {} - - if sso_db_record and sso_db_record.sso_settings: - # Load settings from database - sso_settings_dict = dict(sso_db_record.sso_settings) - - role_mappings_data = sso_settings_dict.pop("role_mappings", None) - role_mappings = None - if role_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - - if isinstance(role_mappings_data, dict): - role_mappings = RoleMappings(**role_mappings_data) - elif isinstance(role_mappings_data, RoleMappings): - role_mappings = role_mappings_data - - team_mappings_data = sso_settings_dict.pop("team_mappings", None) - team_mappings = None - if team_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import TeamMappings - - if isinstance(team_mappings_data, dict): - team_mappings = TeamMappings(**team_mappings_data) - elif isinstance(team_mappings_data, TeamMappings): - team_mappings = team_mappings_data - - decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=sso_settings_dict - ) - - # Build SSO config with database values or environment fallback - - sso_config = SSOConfig( - google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), - google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), - microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None), - microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None), - microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None), - generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None), - generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None), - generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None), - generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None), - generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None), - proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), - user_email=decrypted_sso_settings_dict.get("user_email"), - ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), - role_mappings=role_mappings, - team_mappings=team_mappings, - ) + sso_db_settings = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None + resolved = resolve_sso_config(sso_db_settings, os.environ) # Get the schema for UI display from pydantic import TypeAdapter @@ -826,11 +775,12 @@ async def get_sso_settings(): # Convert to dict for response, masking OAuth client secrets so plaintext # is never sent to the UI. - sso_dict = mask_sensitive_keys(sso_config.model_dump(), _SSO_SENSITIVE_FIELDS) + sso_dict = mask_sensitive_keys(resolved.config.model_dump(), set(SSO_SECRET_FIELDS)) # Add descriptions to the response result = { "values": sso_dict, + "provenance": resolved.provenance, "field_schema": { "description": schema.get("description", ""), "properties": {}, @@ -881,21 +831,6 @@ async def update_sso_settings( detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) - # Update environment variables - env_var_mapping = { - "google_client_id": "GOOGLE_CLIENT_ID", - "google_client_secret": "GOOGLE_CLIENT_SECRET", - "microsoft_client_id": "MICROSOFT_CLIENT_ID", - "microsoft_client_secret": "MICROSOFT_CLIENT_SECRET", - "microsoft_tenant": "MICROSOFT_TENANT", - "generic_client_id": "GENERIC_CLIENT_ID", - "generic_client_secret": "GENERIC_CLIENT_SECRET", - "generic_authorization_endpoint": "GENERIC_AUTHORIZATION_ENDPOINT", - "generic_token_endpoint": "GENERIC_TOKEN_ENDPOINT", - "generic_userinfo_endpoint": "GENERIC_USERINFO_ENDPOINT", - "proxy_base_url": "PROXY_BASE_URL", - } - # Read the existing SSO row first so the audit log captures a real # before/after diff. Stored values are encrypted; decrypt them so the # before-snapshot has the same shape as after_value, and rely on @@ -924,8 +859,8 @@ async def update_sso_settings( # Update environment variables in config and in memory sso_data = sso_config.model_dump() for field_name, value in sso_data.items(): - if field_name in env_var_mapping: - env_var_name = env_var_mapping[field_name] + if field_name in SSO_FIELD_ENV_VARS: + env_var_name = SSO_FIELD_ENV_VARS[field_name] if value: os.environ[env_var_name] = value else: @@ -975,7 +910,7 @@ async def update_sso_settings( else: environment_variables = {} - env_vars_to_remove = set(env_var_mapping.values()) + env_vars_to_remove = set(SSO_FIELD_ENV_VARS.values()) filtered_env_vars = { key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 7234cc2650f..742e0f7818f 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -148,6 +148,10 @@ class SSOConfig(LiteLLMPydanticObjectBase): default=None, description="User info endpoint URL for generic OAuth provider", ) + generic_scope: Optional[str] = Field( + default=None, + description="Space-separated OAuth scopes requested from the generic provider, e.g. 'openid email profile'", + ) # Common settings proxy_base_url: Optional[str] = Field( diff --git a/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py new file mode 100644 index 00000000000..20bea98351f --- /dev/null +++ b/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py @@ -0,0 +1,105 @@ +import os + +from litellm.proxy.config_resolvers._descriptors import FieldDescriptor, resolve_fields +from litellm.proxy.config_resolvers.sso import ( + SSO_FIELD_ENV_VARS, + SSO_SECRET_FIELDS, + resolve_sso_config, +) + +_D = ( + FieldDescriptor("client_id", "client_id", "CLIENT_ID"), + FieldDescriptor("scope", "scope", "SCOPE", default="openid"), +) + + +def test_resolve_fields_db_wins_over_env(): + values, provenance = resolve_fields(_D, {"client_id": "from-db"}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-db" + assert provenance["client_id"] == "db" + + +def test_resolve_fields_blank_db_falls_back_to_env(): + values, provenance = resolve_fields(_D, {"client_id": " "}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-env" + assert provenance["client_id"] == "env" + + +def test_resolve_fields_blank_everywhere_falls_to_default(): + values, provenance = resolve_fields(_D, {}, {"SCOPE": ""}) + assert values["scope"] == "openid" + assert provenance["scope"] == "default" + + +def test_resolve_fields_unset_everywhere(): + values, provenance = resolve_fields(_D, {}, {}) + assert values["client_id"] is None + assert provenance["client_id"] == "unset" + + +def test_resolve_fields_empty_db_absent_by_default_falls_to_env(): + # SSO semantics: a present-but-empty stored value is absent, so env wins. + values, provenance = resolve_fields(_D, {"client_id": ""}, {"CLIENT_ID": "from-env"}) + assert values["client_id"] == "from-env" + assert provenance["client_id"] == "env" + + +def test_resolve_fields_empty_db_is_explicit_clear_when_flag_set(): + # Alerting semantics: a present-but-empty stored value is an explicit clear + # that must win over a stale env var. + values, provenance = resolve_fields( + _D, {"client_id": ""}, {"CLIENT_ID": "stale-env"}, empty_db_is_set=True + ) + assert values["client_id"] == "" + assert provenance["client_id"] == "db" + + +def test_sso_descriptor_mapping_is_single_sourced(): + # The write path and read path both consume this mapping; it must cover every + # env-backed SSO field and map to the uppercase env var. + assert SSO_FIELD_ENV_VARS["generic_client_id"] == "GENERIC_CLIENT_ID" + assert SSO_SECRET_FIELDS == frozenset( + {"google_client_secret", "microsoft_client_secret", "generic_client_secret"} + ) + + +def test_resolve_sso_config_returns_unmasked_secret_and_provenance(): + # The resolver hands back plaintext; masking is the endpoint's job. If the + # resolver masked, the login path would consume a masked secret and fail. + resolved = resolve_sso_config( + {"generic_client_secret": "super-secret-value"}, + {"GENERIC_CLIENT_ID": "env-id"}, + ) + assert resolved.config.generic_client_secret == "super-secret-value" + assert resolved.provenance["generic_client_secret"] == "db" + assert resolved.config.generic_client_id == "env-id" + assert resolved.provenance["generic_client_id"] == "env" + + +def test_resolve_sso_config_parses_structured_mappings(): + resolved = resolve_sso_config( + { + "generic_client_id": "id", + "role_mappings": { + "provider": "generic", + "group_claim": "groups", + "default_role": "internal_user", + "roles": {}, + }, + "team_mappings": {"team_ids_jwt_field": "teams"}, + }, + {}, + ) + assert resolved.config.role_mappings is not None + assert resolved.config.role_mappings.group_claim == "groups" + assert resolved.config.team_mappings is not None + assert resolved.config.team_mappings.team_ids_jwt_field == "teams" + + +def test_resolve_sso_config_does_not_mutate_os_environ(monkeypatch): + # Unlike the legacy read path, resolving must not write os.environ. + monkeypatch.delenv("GENERIC_CLIENT_ID", raising=False) + before = dict(os.environ) + resolve_sso_config({"generic_client_id": "id-from-db"}, os.environ) + assert dict(os.environ) == before + assert "GENERIC_CLIENT_ID" not in os.environ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5effa3073ee..5b67780dc58 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1174,6 +1174,113 @@ def test_get_config_returns_email_settings(monkeypatch): assert "*" in variables["SMTP_PASSWORD"] +def _get_email_alert_variables(monkeypatch, config_data): + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234" + ) + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + email_alert = next((a for a in response.json()["alerts"] if a["name"] == "email"), None) + assert email_alert is not None + return email_alert["variables"] + + +def test_get_config_returns_email_settings_set_only_in_process_env(monkeypatch): + """ + Regression for LIT-4165. + + SMTP supplied purely as process env vars (helm/terraform, no UI writes) is + live at runtime because litellm/proxy/utils.py::send_email resolves every + field from os.getenv. The /get/config/callbacks email block only read the + config/DB environment_variables overlay though, so those deployments saw an + empty Email Server Settings page and could not tell SMTP was configured. + The slack block one branch above already fell back to os.getenv. + """ + smtp_password = "env-only-app-password" + monkeypatch.setenv("SMTP_HOST", "smtp.env-host.com") + monkeypatch.setenv("SMTP_PORT", "2525") + monkeypatch.setenv("SMTP_TLS", "False") + monkeypatch.setenv("SMTP_USERNAME", "env-user") + monkeypatch.setenv("SMTP_PASSWORD", smtp_password) + monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@env-host.com") + monkeypatch.setenv("TEST_EMAIL_ADDRESS", "admin@env-host.com") + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {}, + }, + ) + + # Every one of these was None before the fix, despite SMTP working. + assert variables["SMTP_HOST"] == "smtp.env-host.com" + assert variables["SMTP_PORT"] == "2525" + assert variables["SMTP_TLS"] == "False" + assert variables["SMTP_USERNAME"] == "env-user" + assert variables["SMTP_SENDER_EMAIL"] == "alerts@env-host.com" + assert variables["TEST_EMAIL_ADDRESS"] == "admin@env-host.com" + + # An env-sourced secret is masked exactly like a stored one. + assert variables["SMTP_PASSWORD"] not in (None, smtp_password) + assert "*" in variables["SMTP_PASSWORD"] + + +def test_get_config_email_settings_prefer_stored_over_process_env(monkeypatch): + """ + Stored environment_variables win over the process environment, matching the + load order in ProxyConfig.get_config, which pushes stored values into + os.environ. Only a field with no stored entry falls back to os.getenv. + """ + monkeypatch.setenv("SMTP_HOST", "smtp.env-host.com") + monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@env-host.com") + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {"SMTP_HOST": "smtp.stored-host.com"}, + }, + ) + + assert variables["SMTP_HOST"] == "smtp.stored-host.com" + assert variables["SMTP_SENDER_EMAIL"] == "alerts@env-host.com" + + +def test_get_config_email_settings_absent_everywhere_stay_none(monkeypatch): + """A field set in neither source is reported unset rather than invented.""" + for var in ("SMTP_HOST", "SMTP_PORT", "SMTP_TLS", "SMTP_USERNAME", "SMTP_PASSWORD", "SMTP_SENDER_EMAIL"): + monkeypatch.delenv(var, raising=False) + + variables = _get_email_alert_variables( + monkeypatch, + { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": {}, + }, + ) + + assert variables["SMTP_HOST"] is None + assert variables["SMTP_PASSWORD"] is None + + def test_get_config_returns_slack_webhook(monkeypatch): """ Same double-decryption regression as the email block (issue #19221): the 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 805baed9e1e..85dbf70b452 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 @@ -396,6 +396,146 @@ class TestProxySettingEndpoints: call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args assert call_args.kwargs["where"]["id"] == "sso_config" + def _mock_sso_db_record(self, monkeypatch, sso_settings): + """Point /get/sso_settings at a stored SSO row (or None for no row).""" + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + if sso_settings is None: + mock_db_record = None + else: + mock_db_record = MagicMock() + mock_db_record.sso_settings = sso_settings + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # The resolver decrypts stored values via decrypt_value_helper; make it an + # identity so the plaintext fixtures round-trip. + monkeypatch.setattr( + "litellm.proxy.config_resolvers.sso.decrypt_value_helper", + lambda value, key, exception_type="error", return_original_value=False: value, + ) + + def test_get_sso_settings_falls_back_to_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """ + Regression for LIT-4165. + + SSO configured purely as process env vars (helm/terraform, no UI writes) + logs users in successfully, because ui_sso.py resolves every setting from + os.environ. /get/sso_settings read only the sso_config table though, so + the Admin UI showed "not configured" for a working SSO deployment and hid + the Edit/Delete controls behind an empty-state placeholder. + """ + self._mock_sso_db_record(monkeypatch, None) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "env-client-secret-value") + monkeypatch.setenv("GENERIC_AUTHORIZATION_ENDPOINT", "https://idp.example.com/authorize") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + monkeypatch.setenv("GENERIC_USERINFO_ENDPOINT", "https://idp.example.com/userinfo") + monkeypatch.setenv("GENERIC_SCOPE", "openid email profile groups") + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + + # Every one of these was None before the fix, despite SSO working. + assert values["generic_client_id"] == "env-client-id" + assert values["generic_authorization_endpoint"] == "https://idp.example.com/authorize" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + assert values["generic_userinfo_endpoint"] == "https://idp.example.com/userinfo" + assert values["generic_scope"] == "openid email profile groups" + assert values["proxy_base_url"] == "https://gateway.example.com" + + # An env-sourced secret is masked exactly like a stored one. + assert values["generic_client_secret"] not in (None, "env-client-secret-value") + assert "*" in values["generic_client_secret"] + + def test_get_sso_settings_does_not_mutate_os_environ( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A GET must not write os.environ. The legacy read path decrypted DB + values straight into the environment, so opening the settings page + repopulated env and masked any consumer that stopped reading it.""" + self._mock_sso_db_record(monkeypatch, {"generic_client_id": "db-only-id"}) + monkeypatch.delenv("GENERIC_CLIENT_ID", raising=False) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + assert response.json()["values"]["generic_client_id"] == "db-only-id" + # The DB value must NOT have leaked into the process environment. + assert "GENERIC_CLIENT_ID" not in os.environ + + def test_get_sso_settings_prefers_stored_over_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A stored value wins; only fields absent from the row fall back to env.""" + self._mock_sso_db_record(monkeypatch, {"generic_client_id": "stored-client-id"}) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["generic_client_id"] == "stored-client-id" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + + def test_get_sso_settings_blank_stored_value_falls_back_to_process_env( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """ + Blank means absent. update_sso_settings clears the env var for a blank + field, so a blank row entry cannot describe a live setting; os.environ is + the effective config and is what the UI must report. + """ + self._mock_sso_db_record(monkeypatch, {"generic_client_id": " ", "generic_token_endpoint": ""}) + monkeypatch.setenv("GENERIC_CLIENT_ID", "env-client-id") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + values = response.json()["values"] + assert values["generic_client_id"] == "env-client-id" + assert values["generic_token_endpoint"] == "https://idp.example.com/token" + + def test_get_sso_settings_unset_everywhere_reports_source( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """A field set in neither source is unset (or its effective default), + and provenance reports which.""" + self._mock_sso_db_record(monkeypatch, None) + for env_var in ( + "GENERIC_CLIENT_ID", + "GENERIC_CLIENT_SECRET", + "GENERIC_TOKEN_ENDPOINT", + "GENERIC_SCOPE", + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "PROXY_BASE_URL", + ): + monkeypatch.delenv(env_var, raising=False) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + body = response.json() + values = body["values"] + provenance = body["provenance"] + assert values["generic_client_id"] is None + assert provenance["generic_client_id"] == "unset" + assert values["generic_client_secret"] is None + assert values["google_client_id"] is None + # generic_scope carries the same effective default the login path applies, + # so the settings page shows the scope logins would actually request. + assert values["generic_scope"] == "openid email profile" + assert provenance["generic_scope"] == "default" + def test_update_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating the SSO settings to the dedicated database table""" import json @@ -1463,19 +1603,20 @@ class TestProxySettingEndpoints: monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - # Mock the decryption method to return decrypted values - def mock_decrypt_and_set(environment_variables): - return { - "google_client_id": "decrypted_google_id", - "google_client_secret": "decrypted_google_secret", - "microsoft_client_id": "decrypted_microsoft_id", - "proxy_base_url": "https://decrypted.example.com", - } + # The resolver decrypts each stored value via decrypt_value_helper; map + # the ciphertext fixtures to their plaintext. + decrypted_by_ciphertext = { + "encrypted_google_id": "decrypted_google_id", + "encrypted_google_secret": "decrypted_google_secret", + "encrypted_microsoft_id": "decrypted_microsoft_id", + "encrypted_proxy_url": "https://decrypted.example.com", + } - from litellm.proxy.proxy_server import proxy_config + def mock_decrypt(value, key, exception_type="error", return_original_value=False): + return decrypted_by_ciphertext.get(value, value) monkeypatch.setattr( - proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set + "litellm.proxy.config_resolvers.sso.decrypt_value_helper", mock_decrypt ) response = client.get("/get/sso_settings") diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts index 0431a8d39f7..1a02e363de9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts @@ -24,6 +24,7 @@ export interface SSOSettingsValues { generic_authorization_endpoint: string | null; generic_token_endpoint: string | null; generic_userinfo_endpoint: string | null; + generic_scope: string | null; proxy_base_url: string | null; user_email: string | null; ui_access_mode: string | null; diff --git a/ui/litellm-dashboard/src/components/SSOModals.test.tsx b/ui/litellm-dashboard/src/components/SSOModals.test.tsx index 792a964f01c..c5da4ee7064 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.test.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.test.tsx @@ -462,6 +462,7 @@ describe("SSOModals", () => { generic_authorization_endpoint: null, generic_token_endpoint: null, generic_userinfo_endpoint: null, + generic_scope: null, proxy_base_url: null, user_email: null, sso_provider: null, diff --git a/ui/litellm-dashboard/src/components/SSOModals.tsx b/ui/litellm-dashboard/src/components/SSOModals.tsx index 88ce72573d7..637abbf4a81 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.tsx @@ -1,11 +1,12 @@ import React, { useEffect, useState } from "react"; -import { Modal, Form, Input, Button as Button2, Select, Checkbox } from "antd"; +import { Modal, Form, Button as Button2, Select, Checkbox } from "antd"; import { Text, TextInput } from "@tremor/react"; import { getSSOSettings, updateSSOSettings } from "./networking"; import NotificationsManager from "./molecules/notifications_manager"; import { parseErrorMessage } from "./shared/errorUtils"; import { Logo } from "@/components/molecules/logo/Logo"; import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./Settings/AdminSettings/SSOSettings/constants"; +import { renderProviderFields } from "./Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm"; interface SSOModalsProps { isAddSSOModalVisible: boolean; @@ -20,82 +21,6 @@ interface SSOModalsProps { ssoConfigured?: boolean; // Add optional prop to indicate if SSO is configured } -// Define the SSO provider configuration type -interface SSOProviderConfig { - envVarMap: Record; - fields: Array<{ - label: string; - name: string; - placeholder?: string; - }>; -} - -// Define configurations for each SSO provider -const ssoProviderConfigs: Record = { - google: { - envVarMap: { - google_client_id: "GOOGLE_CLIENT_ID", - google_client_secret: "GOOGLE_CLIENT_SECRET", - }, - fields: [ - { label: "Google Client ID", name: "google_client_id" }, - { label: "Google Client Secret", name: "google_client_secret" }, - ], - }, - microsoft: { - envVarMap: { - microsoft_client_id: "MICROSOFT_CLIENT_ID", - microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", - microsoft_tenant: "MICROSOFT_TENANT", - }, - fields: [ - { label: "Microsoft Client ID", name: "microsoft_client_id" }, - { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, - { label: "Microsoft Tenant", name: "microsoft_tenant" }, - ], - }, - okta: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { - label: "Authorization Endpoint", - name: "generic_authorization_endpoint", - placeholder: "https://your-domain/authorize", - }, - { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, - { - label: "Userinfo Endpoint", - name: "generic_userinfo_endpoint", - placeholder: "https://your-domain/userinfo", - }, - ], - }, - generic: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, - { label: "Token Endpoint", name: "generic_token_endpoint" }, - { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, - ], - }, -}; - const SSOModals: React.FC = ({ isAddSSOModalVisible, isInstructionsModalVisible, @@ -266,6 +191,7 @@ const SSOModals: React.FC = ({ generic_authorization_endpoint: null, generic_token_endpoint: null, generic_userinfo_endpoint: null, + generic_scope: null, proxy_base_url: null, user_email: null, sso_provider: null, @@ -291,22 +217,6 @@ const SSOModals: React.FC = ({ }; // Helper function to render provider fields - const renderProviderFields = (provider: string) => { - const config = ssoProviderConfigs[provider]; - if (!config) return null; - - return config.fields.map((field) => ( - - {field.name.includes("client") ? : } - - )); - }; - return ( <> = ({ ); }; -export { ssoProviderConfigs }; // Export for use in other components export default SSOModals; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx index 21132fff63b..2f9c49dfa56 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx @@ -2,7 +2,7 @@ import { Form } from "antd"; import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import { renderWithProviders } from "../../../../../../tests/test-utils"; import { afterEach, describe, expect, it, vi } from "vitest"; -import BaseSSOSettingsForm, { renderProviderFields } from "./BaseSSOSettingsForm"; +import BaseSSOSettingsForm, { renderProviderFields, ssoProviderConfigs } from "./BaseSSOSettingsForm"; describe("BaseSSOSettingsForm", () => { afterEach(() => { @@ -285,13 +285,70 @@ describe("renderProviderFields", () => { it("should return fields for okta provider", () => { const result = renderProviderFields("okta"); expect(result).not.toBeNull(); - expect(result?.length).toBe(5); + expect(result?.length).toBe(6); }); it("should return fields for generic provider", () => { const result = renderProviderFields("generic"); expect(result).not.toBeNull(); - expect(result?.length).toBe(5); + expect(result?.length).toBe(6); + }); + + it.each(["okta", "generic"])( + "renders an optional generic_scope field for %s so editing cannot clear it", + (provider) => { + const scopeField = ssoProviderConfigs[provider].fields.find((field) => field.name === "generic_scope"); + expect(scopeField).toBeDefined(); + expect(scopeField?.required).toBe(false); + expect(ssoProviderConfigs[provider].envVarMap.generic_scope).toBe("GENERIC_SCOPE"); + }, + ); + + it("submits generic_scope untouched, so saving an unrelated edit cannot clear GENERIC_SCOPE", async () => { + // update_sso_settings clears the env var for any mapped field its payload + // omits, and antd only submits mounted fields. So the Scopes field being + // present is what stops an unrelated edit from downgrading a custom scope + // to the provider default. Dropping the field from ssoProviderConfigs must + // fail here rather than silently in production. + const handleSubmit = vi.fn(); + let form: any; + const TestWrapper = () => { + const [formInstance] = Form.useForm(); + form = formInstance; + return ; + }; + + renderWithProviders(); + + // Mirror EditSSOSettingsModal hydrating the form from the GET response. + await act(async () => { + form.setFieldsValue({ + sso_provider: "generic", + generic_client_id: "client-id", + generic_client_secret: "client-secret", + generic_authorization_endpoint: "https://idp.example.com/authorize", + generic_token_endpoint: "https://idp.example.com/token", + generic_userinfo_endpoint: "https://idp.example.com/userinfo", + generic_scope: "openid email profile groups", + proxy_base_url: "https://gateway.example.com", + user_email: "admin@example.com", + }); + }); + + // The admin edits something else entirely and saves. + await act(async () => { + form.setFieldsValue({ generic_token_endpoint: "https://idp.example.com/token/v2" }); + form.submit(); + }); + + await waitFor(() => { + expect(handleSubmit).toHaveBeenCalledWith( + expect.objectContaining({ + generic_token_endpoint: "https://idp.example.com/token/v2", + generic_scope: "openid email profile groups", + }), + ); + }); }); it("renders provider logos in the dropdown and falls back to a letter avatar on load error", async () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx index 6971c107a73..caa6ff4f1e8 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -18,6 +18,7 @@ export interface SSOProviderConfig { label: string; name: string; placeholder?: string; + required?: boolean; }>; } @@ -52,6 +53,7 @@ export const ssoProviderConfigs: Record = { generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + generic_scope: "GENERIC_SCOPE", }, fields: [ { label: "Generic Client ID", name: "generic_client_id" }, @@ -67,6 +69,7 @@ export const ssoProviderConfigs: Record = { name: "generic_userinfo_endpoint", placeholder: "https://your-domain/userinfo", }, + { label: "Scopes", name: "generic_scope", placeholder: "openid email profile", required: false }, ], }, generic: { @@ -76,6 +79,7 @@ export const ssoProviderConfigs: Record = { generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + generic_scope: "GENERIC_SCOPE", }, fields: [ { label: "Generic Client ID", name: "generic_client_id" }, @@ -83,6 +87,7 @@ export const ssoProviderConfigs: Record = { { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, { label: "Token Endpoint", name: "generic_token_endpoint" }, { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, + { label: "Scopes", name: "generic_scope", placeholder: "openid email profile", required: false }, ], }, }; @@ -97,7 +102,7 @@ export const renderProviderFields = (provider: string) => { key={field.name} label={field.label} name={field.name} - rules={[{ required: true, message: `Please enter the ${field.label.toLowerCase()}` }]} + rules={[{ required: field.required !== false, message: `Please enter the ${field.label.toLowerCase()}` }]} > {field.name.includes("client") ? : } diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index e3361050422..849921c1076 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -111,6 +111,7 @@ export default function SSOSettings() { label: "User Info Endpoint", render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint), }, + { label: "Scopes", render: (values: SSOSettingsValues) => renderSimpleValue(values.generic_scope) }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, isTeamMappingsEnabled ? { @@ -143,6 +144,7 @@ export default function SSOSettings() { label: "User Info Endpoint", render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint), }, + { label: "Scopes", render: (values: SSOSettingsValues) => renderSimpleValue(values.generic_scope) }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, isTeamMappingsEnabled ? { diff --git a/ui/litellm-dashboard/src/components/email_settings.tsx b/ui/litellm-dashboard/src/components/email_settings.tsx index 83968a458b9..2a7988832ea 100644 --- a/ui/litellm-dashboard/src/components/email_settings.tsx +++ b/ui/litellm-dashboard/src/components/email_settings.tsx @@ -26,9 +26,17 @@ const EmailSettings: React.FC = ({ accessToken, premiumUser, .forEach((alert) => { Object.entries(alert.variables ?? {}).forEach(([key, value]) => { const inputElement = document.querySelector(`input[name="${key}"]`) as HTMLInputElement; - if (inputElement && inputElement.value) { - updatedVariables[key] = inputElement?.value; + if (!inputElement || !inputElement.value) { + return; } + // Only send fields the admin actually edited. Values rendered from the + // server are masked (SMTP_PASSWORD) or sourced from the process + // environment, so re-submitting an untouched field would persist a mask + // or copy env-managed config into the database. + if (inputElement.value === (value == null ? "" : String(value))) { + return; + } + updatedVariables[key] = inputElement.value; }); }); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 203057f23f1..d109fc1783d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30686,6 +30686,11 @@ export interface components { * @description Generic OAuth Client Secret for SSO authentication */ generic_client_secret?: string | null; + /** + * Generic Scope + * @description Space-separated OAuth scopes requested from the generic provider, e.g. 'openid email profile' + */ + generic_scope?: string | null; /** * Generic Token Endpoint * @description Token endpoint URL for generic OAuth provider @@ -30750,6 +30755,10 @@ export interface components { field_schema: { [key: string]: unknown; }; + /** Provenance */ + provenance?: { + [key: string]: string; + }; /** Values */ values: { [key: string]: unknown; From 8d217a4d5f9f5511fc6068a61e5718d42a2b8200 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 22 Jul 2026 15:15:37 -0700 Subject: [PATCH 043/130] fix(scim): parse membership id from filtered PATCH path when value omitted (#34181) Okta commonly sends SCIM membership removals as a filtered path with no request body value, e.g. Groups PATCH members[value eq "uid"] and Users PATCH groups[value eq "tid"]. The patch handlers pulled ids only from op.value, so these removes were a silent no-op and the member or team was never dropped Add a linear-time filter parser reused by both the Groups members path and the Users groups path so the id is taken from the [value eq "..."] filter when op.value is absent, for add and remove ops. The eq operator is matched case-insensitively, both quote styles are accepted, and the quoted value is unescaped. The path-filter fallback only fires when the request body value is omitted, so an explicit empty value no longer resurrects the filter id, and the compared value must be quoted per the SCIM filter grammar --- .../management_endpoints/scim/scim_v2.py | 33 +++- .../scim/test_scim_patch_user.py | 55 +++++++ .../scim/test_scim_v2_endpoints.py | 142 ++++++++++++++++++ 3 files changed, 228 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 525ce9f0e89..90eae5bbb21 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1353,6 +1353,31 @@ def _extract_group_values(value: Any) -> List[str]: return group_values +def _extract_ids_from_path_filter(path: str | None, attribute: str) -> List[str]: + """Return ids from a SCIM filtered path like ``members[value eq "id"]``. + + Okta commonly sends membership removals as a filtered path and omits the + request body ``value``, so the id lives only inside the ``[value eq "..."]`` + filter. The ``eq`` operator is matched case-insensitively per the SCIM + spec; the id keeps its original case. Per the SCIM filter grammar the + compared value must be quoted (single or double), so malformed unquoted + filters yield no id. A quoted id may contain escaped quotes and + backslashes (``\\"`` and ``\\\\``), which are unescaped before use. + ``path`` must be the raw, case-preserving path from the patch op. + """ + if not path: + return [] + match = re.match( + rf"""\s*{re.escape(attribute)}\s*\[\s*value\s+eq\s+(['"])((?:\\.|[^\\])*?)\1\s*\]\s*$""", + path, + flags=re.IGNORECASE, + ) + if not match: + return [] + extracted = re.sub(r"\\(.)", r"\1", match.group(2)) + return [extracted] if extracted else [] + + def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: """Handle displayname updates.""" if op_type == "remove": @@ -1396,9 +1421,11 @@ def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict scim_metadata["familyName"] = str(value) -def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> Optional[Set[str]]: +def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str], path: str | None) -> Set[str] | None: """Handle group/team membership operations.""" group_values = _extract_group_values(value) + if not group_values and value is None: + group_values = _extract_ids_from_path_filter(path, "groups") if op_type == "replace": return set(group_values) elif op_type == "add": @@ -1511,7 +1538,7 @@ def _apply_patch_ops( elif _multi_valued_attribute_base(path) in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS: _handle_multi_valued_attribute_update(path, op_type, value, metadata) elif path.startswith("groups"): - new_replace_set = _handle_group_operations(op_type, value, teams_set) + new_replace_set = _handle_group_operations(op_type, value, teams_set, op.path) if new_replace_set is not None: replace_team_set = new_replace_set else: @@ -1975,6 +2002,8 @@ async def _process_group_patch_operations( elif path.startswith("members"): # Handle member operations member_values = _extract_group_values(value) + if not member_values and value is None: + member_values = _extract_ids_from_path_filter(op.path, "members") # Check the feature flag scim_upsert_user = await _get_scim_upsert_user_setting() # Validate all users exist or create them based on feature flag diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py index f8995a6f4da..c3af4208d37 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_patch_user.py @@ -465,3 +465,58 @@ def test_apply_patch_ops_filtered_path_raises_400_instead_of_junk_metadata(): ) assert exc_info.value.status_code == 400 + + +def test_apply_patch_ops_remove_group_filtered_path_without_value(): + """Okta removes a user from a team with groups[value eq "..."] and no body + value; the team id must be parsed from the filter so the remove takes effect""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1", "team-2"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="remove", path='groups[value eq "team-1"]')] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == {"team-2"} + + +def test_apply_patch_ops_add_group_filtered_path_without_value(): + """A filtered add path with no body value adds the team id from the filter.""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="add", path="groups[value eq 'team-3']")] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == {"team-1", "team-3"} + + +def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter(): + """A filtered replace with an explicit empty value must not resurrect the + filter id; the team set is replaced with the empty value as given.""" + user = LiteLLM_UserTable( + user_id="user-fp", + user_email="fp@example.com", + teams=["team-1", "team-2"], + metadata={}, + ) + patch_ops = SCIMPatchOp( + Operations=[ + SCIMPatchOperation(op="replace", path='groups[value eq "team-1"]', value=[]) + ] + ) + + _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) + + assert final_team_set == set() diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 850b8fc7cfd..7bb74285ac6 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,3 +1,4 @@ +import time from unittest.mock import AsyncMock import pytest @@ -17,6 +18,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( UserProvisionerHelpers, _apply_group_patch_updates, _extract_group_member_ids, + _extract_ids_from_path_filter, _handle_team_membership_changes, _process_group_patch_operations, _recompute_scim_member_roles, @@ -3377,3 +3379,143 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock call.kwargs["user_id"] for call in calls if call.kwargs.get("teams_ids_to_add_user_to") == [group_id] } assert added_user_ids == set() + + +@pytest.mark.parametrize( + "path, attribute, expected", + [ + ('members[value eq "user-1"]', "members", ["user-1"]), + ("members[value eq 'user-1']", "members", ["user-1"]), + ('members[value EQ "user-1"]', "members", ["user-1"]), + ('members[ value eq "user-1" ]', "members", ["user-1"]), + ('groups[value eq "team-1"]', "groups", ["team-1"]), + ('members[value eq "Mixed-CASE-Id"]', "members", ["Mixed-CASE-Id"]), + ('members[value eq "a\\"b"]', "members", ['a"b']), + ('members[value eq "a\\\\b"]', "members", ["a\\b"]), + ("members[value eq 'a\\'b']", "members", ["a'b"]), + ("members", "members", []), + ('groups[value eq "team-1"]', "members", []), + (None, "members", []), + ('members[value eq ""]', "members", []), + ("members[value eq user-1]", "members", []), + ("members[value eq unintendeduser]", "members", []), + ], +) +def test_extract_ids_from_path_filter(path, attribute, expected): + assert _extract_ids_from_path_filter(path, attribute) == expected + + +def test_extract_ids_from_path_filter_unterminated_is_linear(): + """A pathological unterminated quoted filter must not trigger super-linear + backtracking; it returns no id and completes near-instantly.""" + pathological = 'members[value eq "' + ("\\" * 200) + + start = time.perf_counter() + result = _extract_ids_from_path_filter(pathological, "members") + elapsed = time.perf_counter() - start + + assert result == [] + assert elapsed < 1.0 + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_filtered_path_without_value(mocker): + """Okta sends group membership removals as a filtered path with no request + body value; the member id must be parsed out of members[value eq "..."]""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path='members[value eq "user-1"]')], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="user-1", role="user"), + Member(user_id="user-2", role="user"), + ], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-1") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == {"user-2"} + + +@pytest.mark.asyncio +async def test_process_group_patch_add_filtered_path_without_value(mocker): + """A filtered add path with no body value adds the id parsed from the filter.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="add", path='members[value eq "user-3"]')], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[Member(user_id="user-1", role="user")], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-3") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == {"user-1", "user-3"} + + +@pytest.mark.asyncio +async def test_process_group_patch_replace_empty_value_does_not_use_path_filter(mocker): + """An explicit empty replace value must clear membership rather than pull an + id from the filtered path, which would retain one member and drop the rest.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="replace", path='members[value eq "user-1"]', value=[]) + ], + ) + + existing_team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="Team One", + members=[], + members_with_roles=[ + Member(user_id="user-1", role="user"), + Member(user_id="user-2", role="user"), + ], + ) + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="user-1") + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=prisma_client, + ) + + assert final_members == set() From 0b4446edccf282cc63a29d7f6b8415211b13fc3f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:16:52 -0700 Subject: [PATCH 044/130] fix(ui): remove misleading os.environ tooltip from logging settings (#34305) Co-authored-by: yuneng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/components/team/LoggingSettings.tsx | 3 --- 1 file changed, 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx index 3572082872e..80bea3ef2d2 100644 --- a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx +++ b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx @@ -117,9 +117,6 @@ const LoggingSettings: React.FC = ({
= ({ accessToken, userRole }) => { /> {promptToDelete && ( - { + if (!open && !isDeleting) handleDeleteCancel(); + }} > -

Are you sure you want to delete prompt: {promptToDelete.name} ?

-

This action cannot be undone.

- + + + Delete Prompt + + Are you sure you want to delete prompt: {promptToDelete.name} ? This action cannot be undone. + + + + Cancel + + + + )} ); From 692e6d48e983832939d7e3e26c911e21c9ee7ef3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 22 Jul 2026 17:01:13 -0700 Subject: [PATCH 053/130] refactor(ui): migrate old-usage to shadcn (#34304) * test(ui): characterise the old usage page before migrating it Role- and text-based coverage of the route as it behaves on Tremor, so the shadcn migration has a regression net it did not get to write. Pins the DISABLE_EXPENSIVE_DB_QUERIES branch (warning copy, the docs link and its target, and that every expensive query is skipped), the admin vs non-admin tab set, the cost cards, and the provider and customer tables * refactor(ui): migrate old-usage to shadcn Replaces Tremor with the installed shadcn primitives and the shared recharts wrappers on the only file the route owns. Tabs, cards, tables, the key select and the tag multi-select come from src/components/ui; the bar, area and donut charts come from src/components/shared/charts. Tremor BarList has no shared equivalent, so Total Spend Per Team is composed from ui/meter, which also means the per-team totals stay numbers in state instead of pre-formatted strings; a team total of 1,000 or more used to make the bar widths NaN. The Database Query Limit Reached warning moves with it: same copy, same docs link, still short-circuiting every expensive query. Drops the file's no-restricted-imports suppression and the dead customTooltip, getTopKeys, DataDict and UserData symbols. The characterisation test from the previous commit is unchanged and green on both sides --- ui/litellm-dashboard/eslint-suppressions.json | 3 - .../old-usage/_components/usage.test.tsx | 202 +++++ .../old-usage/_components/usage.tsx | 855 +++++++++--------- 3 files changed, 622 insertions(+), 438 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index f8927becba0..b9b27e6cbb5 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -664,9 +664,6 @@ } }, "src/app/(dashboard)/old-usage/_components/usage.tsx": { - "no-restricted-imports": { - "count": 2 - }, "react-hooks/immutability": { "count": 1 }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx new file mode 100644 index 00000000000..e3db50b7300 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx @@ -0,0 +1,202 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../../../tests/test-utils"; +import UsagePage from "./usage"; + +const networking = vi.hoisted(() => ({ + adminSpendLogsCall: vi.fn(), + adminTopKeysCall: vi.fn(), + adminTopModelsCall: vi.fn(), + adminTopEndUsersCall: vi.fn(), + teamSpendLogsCall: vi.fn(), + tagsSpendLogsCall: vi.fn(), + allTagNamesCall: vi.fn(), + adminspendByProvider: vi.fn(), + adminGlobalActivity: vi.fn(), + adminGlobalActivityPerModel: vi.fn(), + getProxyUISettings: vi.fn(), + modelAvailableCall: vi.fn(), + keyInfoV1Call: vi.fn(), +})); + +vi.mock("@/components/networking", () => networking); +vi.mock("../../../../components/networking", () => networking); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: "sk-test", + token: "tok", + userRole: "Admin", + userId: "u1", + premiumUser: true, + }), +})); + +const UNLIMITED_SETTINGS = { DISABLE_EXPENSIVE_DB_QUERIES: false, NUM_SPEND_LOGS_ROWS: 10 }; + +const renderUsage = (overrides: Partial> = {}) => + renderWithProviders( + , + ); + +beforeEach(() => { + vi.clearAllMocks(); + networking.getProxyUISettings.mockResolvedValue(UNLIMITED_SETTINGS); + networking.adminSpendLogsCall.mockResolvedValue([{ date: "2026-07-01", spend: 12.5 }]); + networking.adminTopKeysCall.mockResolvedValue([ + { api_key: "sk-abcdefghijk", key_alias: "prod-key", total_spend: 9.5 }, + ]); + networking.adminTopModelsCall.mockResolvedValue([{ model: "gpt-5.1", total_spend: 7.25 }]); + networking.adminTopEndUsersCall.mockResolvedValue([ + { end_user: "customer-alpha", total_spend: 3.5, total_count: 42 }, + ]); + networking.teamSpendLogsCall.mockResolvedValue({ + daily_spend: [{ date: "2026-07-01", "team-a": 5 }], + teams: ["team-a"], + total_spend_per_team: [{ team_id: "team-a", total_spend: 5 }], + }); + networking.tagsSpendLogsCall.mockResolvedValue({ spend_per_tag: [{ name: "prod", spend: 4 }] }); + networking.allTagNamesCall.mockResolvedValue({ tag_names: ["prod", "staging"] }); + networking.adminspendByProvider.mockResolvedValue([{ provider: "openai", spend: 6.75 }]); + networking.adminGlobalActivity.mockResolvedValue({ + sum_api_requests: 120, + sum_total_tokens: 4500, + daily_data: [{ date: "2026-07-01", api_requests: 120, total_tokens: 4500 }], + }); + networking.adminGlobalActivityPerModel.mockResolvedValue([]); + networking.modelAvailableCall.mockResolvedValue({ data: [] }); + networking.keyInfoV1Call.mockResolvedValue({ info: {} }); +}); + +describe("old usage page", () => { + describe("when the proxy has disabled expensive DB queries", () => { + beforeEach(() => { + networking.getProxyUISettings.mockResolvedValue({ + DISABLE_EXPENSIVE_DB_QUERIES: true, + NUM_SPEND_LOGS_ROWS: 2500000, + }); + }); + + it("shows the database query limit warning instead of the usage dashboard", async () => { + renderUsage(); + + expect(await screen.findByText("Database Query Limit Reached")).toBeInTheDocument(); + expect(screen.getByText(/SpendLogs in DB has/)).toHaveTextContent("2500000"); + expect(screen.getByText(/Please follow our guide to view usage when SpendLogs has more than 1M rows/i)); + expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); + }); + + it("links to the cost tracking guide in a new tab", async () => { + renderUsage(); + + const link = await screen.findByRole("link", { name: "View Usage Guide" }); + expect(link).toHaveAttribute("href", "https://docs.litellm.ai/docs/proxy/cost_tracking"); + expect(link).toHaveAttribute("target", "_blank"); + }); + + it("skips every expensive usage query", async () => { + renderUsage(); + + await screen.findByText("Database Query Limit Reached"); + await waitFor(() => expect(networking.getProxyUISettings).toHaveBeenCalled()); + + expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.adminspendByProvider).not.toHaveBeenCalled(); + expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); + expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); + expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); + expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); + expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); + expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); + }); + }); + + describe("as an admin", () => { + it("renders the admin tabs", async () => { + renderUsage(); + + expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Team Based Usage" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Customer Usage" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Tag Based Usage" })).toBeInTheDocument(); + }); + + it("renders the cost panel cards", async () => { + renderUsage(); + + expect(await screen.findByText("Monthly Spend")).toBeInTheDocument(); + expect(screen.getByText("Top Virtual Keys")).toBeInTheDocument(); + expect(screen.getByText("Top Models")).toBeInTheDocument(); + expect(screen.getByText("Spend by Provider")).toBeInTheDocument(); + }); + + it("lists spend by provider in a table", async () => { + renderUsage(); + + const providerCell = await screen.findByText("openai"); + const row = providerCell.closest("tr"); + expect(row).not.toBeNull(); + expect(within(row as HTMLElement).getByText("$6.75")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument(); + }); + + it("shows the customer usage table when its tab is selected", async () => { + const user = userEvent.setup(); + renderUsage(); + + await user.click(await screen.findByRole("tab", { name: "Customer Usage" })); + + const customerCell = await screen.findByText("customer-alpha"); + const row = customerCell.closest("tr"); + expect(row).not.toBeNull(); + expect(within(row as HTMLElement).getByText("$3.50")).toBeInTheDocument(); + expect(within(row as HTMLElement).getByText("42")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Total Events" })).toBeInTheDocument(); + }); + + it("shows the tag spend panel when its tab is selected", async () => { + const user = userEvent.setup(); + renderUsage(); + + await user.click(await screen.findByRole("tab", { name: "Tag Based Usage" })); + + expect(await screen.findByText("Spend Per Tag")).toBeInTheDocument(); + }); + + it("shows the team spend panel when its tab is selected", async () => { + const user = userEvent.setup(); + renderUsage(); + + await user.click(await screen.findByRole("tab", { name: "Team Based Usage" })); + + expect(await screen.findByText("Total Spend Per Team")).toBeInTheDocument(); + expect(screen.getByText("Daily Spend Per Team")).toBeInTheDocument(); + }); + }); + + describe("as a non-admin", () => { + it("renders only the All Up tab and skips admin-only queries", async () => { + renderUsage({ userRole: "Internal User" }); + + expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Team Based Usage" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Customer Usage" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Tag Based Usage" })).not.toBeInTheDocument(); + + await waitFor(() => expect(networking.adminSpendLogsCall).toHaveBeenCalled()); + expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index 01f8cb1cd45..3d55f9bb698 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -1,40 +1,26 @@ -import { - BarChart, - BarList, - Card, - Title, - Table, - TableHead, - TableHeaderCell, - TableRow, - TableCell, - TableBody, - Subtitle, -} from "@tremor/react"; - import React, { useState, useEffect } from "react"; import ViewUserSpend from "@/components/view_user_spend"; import { ProxySettings } from "@/components/user_dashboard"; import UsageDatePicker from "@/components/shared/usage_date_picker"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { - Grid, - Col, - Text, - TabPanel, - TabPanels, - TabGroup, - TabList, - Tab, - Select, - SelectItem, - DateRangePickerValue, - DonutChart, - AreaChart, - Button, - MultiSelect, - MultiSelectItem, -} from "@tremor/react"; + Combobox, + ComboboxChip, + ComboboxChips, + ComboboxChipsInput, + ComboboxContent, + ComboboxEmpty, + ComboboxItem, + ComboboxList, + ComboboxValue, +} from "@/components/ui/combobox"; +import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { AreaChart, BarChart, DonutChart } from "@/components/shared/charts"; import { adminSpendLogsCall, @@ -68,69 +54,41 @@ interface GlobalActivityData { daily_data: { date: string; api_requests: number; total_tokens: number }[]; } -type CustomTooltipTypeBar = { - payload: any; - active: boolean | undefined; - label: any; -}; +type UsageDateRange = { from?: Date; to?: Date }; -const customTooltip = (props: CustomTooltipTypeBar) => { - const { payload, active } = props; - if (!active || !payload) return null; +type TeamSpendTotal = { name: string; value: number }; - const value = payload[0].payload; - const date = value["startTime"]; - const model_values = value["models"]; - const entries: [string, number][] = Object.entries(model_values).map(([key, value]) => [key, value as number]); +type TagOption = { value: string; label: string; disabled: boolean }; - entries.sort((a, b) => b[1] - a[1]); - const topEntries = entries.slice(0, 5); - - return ( -
- {date} - {topEntries.map(([key, value]) => ( -
-
-

- {key} - {":"} - - {" "} - {value ? `$${formatNumberWithCommas(value, 2)}` : ""} - -

-
-
- ))} -
- ); -}; - -function getTopKeys(data: Array<{ [key: string]: unknown }>): any[] { - const spendKeys: { key: string; spend: unknown }[] = []; - - data.forEach((dict) => { - Object.entries(dict).forEach(([key, value]) => { - if (key !== "spend" && key !== "startTime" && key !== "models" && key !== "users") { - spendKeys.push({ key, spend: value }); - } - }); - }); - - spendKeys.sort((a, b) => Number(b.spend) - Number(a.spend)); - - const topKeys = spendKeys.slice(0, 5).map((k) => k.key); - return topKeys; -} -type DataDict = { [key: string]: unknown }; -type UserData = { user_id: string; spend: number }; +const ALL_TAGS = "all-tags"; const isAdminOrAdminViewer = (role: string | null): boolean => { if (role === null) return false; return role === "Admin" || role === "Admin Viewer"; }; +const TeamSpendBarList: React.FC<{ data: TeamSpendTotal[] }> = ({ data }) => { + const max = Math.max(0, ...data.map((team) => team.value)); + + return ( +
+ {data.map((team) => ( +
+

{team.name}

+ + + + + +

+ {formatNumberWithCommas(team.value, 2)} +

+
+ ))} +
+ ); +}; + const UsagePage: React.FC = ({ accessToken, token, userRole, userID, keys, premiumUser }) => { const currentDate = new Date(); const [keySpendData, setKeySpendData] = useState([]); @@ -141,13 +99,13 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use const [topTagsData, setTopTagsData] = useState([]); const [allTagNames, setAllTagNames] = useState([]); const [uniqueTeamIds, setUniqueTeamIds] = useState([]); - const [totalSpendPerTeam, setTotalSpendPerTeam] = useState([]); + const [totalSpendPerTeam, setTotalSpendPerTeam] = useState([]); const [spendByProvider, setSpendByProvider] = useState([]); const [globalActivity, setGlobalActivity] = useState({} as GlobalActivityData); const [globalActivityPerModel, setGlobalActivityPerModel] = useState([]); - const [selectedKeyID, setSelectedKeyID] = useState(""); - const [selectedTags, setSelectedTags] = useState(["all-tags"]); - const [dateValue, setDateValue] = useState({ + const [selectedKeyToken, setSelectedKeyToken] = useState(null); + const [selectedTags, setSelectedTags] = useState([ALL_TAGS]); + const [dateValue, setDateValue] = useState({ from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), to: new Date(), }); @@ -160,6 +118,21 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use let startTime = formatDate(firstDay); let endTime = formatDate(lastDay); + const selectableKeys: { token: string; alias: string }[] = (keys ?? []) + .filter((key: any) => key && typeof key["key_alias"] === "string" && key["key_alias"].length > 0) + .map((key: any) => ({ token: String(key["token"]), alias: String(key["key_alias"]) })); + + const tagOptions: TagOption[] = [ + { value: ALL_TAGS, label: "All Tags", disabled: false }, + ...allTagNames + .filter((tag) => tag !== ALL_TAGS) + .map((tag) => ({ + value: tag, + label: premiumUser ? tag : `✨ ${tag} (Enterprise only Feature)`, + disabled: !premiumUser, + })), + ]; + function valueFormatterNumbers(number: number) { const formatter = new Intl.NumberFormat("en-US", { maximumFractionDigits: 0, @@ -405,7 +378,7 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use setUniqueTeamIds(teamSpend.teams); return teamSpend.total_spend_per_team.map((tspt: any) => ({ name: tspt["team_id"] || "", - value: formatNumberWithCommas(tspt["total_spend"] || 0, 2), + value: Number(tspt["total_spend"] || 0), })); }, setTotalSpendPerTeam, @@ -524,223 +497,252 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use if (proxySettings?.DISABLE_EXPENSIVE_DB_QUERIES) { return ( -
+
- Database Query Limit Reached - - SpendLogs in DB has {proxySettings.NUM_SPEND_LOGS_ROWS} rows. -

- Please follow our guide to view usage when SpendLogs has more than 1M rows. - - + + Database Query Limit Reached + + +

+ SpendLogs in DB has {proxySettings.NUM_SPEND_LOGS_ROWS} rows. +

+ Please follow our guide to view usage when SpendLogs has more than 1M rows. +

+
); } return ( -
- - - All Up +
+ + + All Up - {isAdminOrAdminViewer(userRole) ? ( + {isAdminOrAdminViewer(userRole) && ( <> - Team Based Usage - Customer Usage - Tag Based Usage - - ) : ( - <> -
+ Team Based Usage + Customer Usage + Tag Based Usage )} - - - - - - Cost - Activity - - - - - - - Project Spend {new Date().toLocaleString("default", { month: "long" })} 1 -{" "} - {new Date(new Date().getFullYear(), new Date().getMonth() + 1, 0).getDate()} - - - - - - Monthly Spend - + + + + + Cost + Activity + + + +
+
+

+ Project Spend {new Date().toLocaleString("default", { month: "long" })} 1 -{" "} + {new Date(new Date().getFullYear(), new Date().getMonth() + 1, 0).getDate()} +

+ +
+
+ + + Monthly Spend + + + + + +
+
+ + + Top Virtual Keys + + + {}} /> + + +
+
+ + + Top Models + + + `$${formatNumberWithCommas(value, 2)}`} + /> + + +
+
+
+ + + Spend by Provider + + +
+
+ `$${formatNumberWithCommas(value, 2)}`} + /> +
+
+ + + + Provider + Spend + + + + {spendByProvider.map((provider) => ( + + {provider.provider} + + + + + ))} + +
+
+
+ + +
+
+ + + +
+ + + All Up + + +
+
+

+ API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} +

+ - - - - - Top Virtual Keys - {}} /> - - - - - Top Models +
+
+

+ Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} +

`$${formatNumberWithCommas(value, 2)}`} + categories={["total_tokens"]} /> - - - - - - Spend by Provider - <> - - - `$${formatNumberWithCommas(value, 2)}`} - /> - - - - - - Provider - Spend - - - - {spendByProvider.map((provider) => ( - - {provider.provider} - - - - - ))} - -
- - - - - - - - - - - All Up - - - +
+
+ + + + {globalActivityPerModel.map((globalActivity, index) => ( + + + {globalActivity.model} + + +
+
+

API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} - +

- - - +
+
+

Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} - +

- - - +
+
+ + + ))} +
+ + + - <> - {globalActivityPerModel.map((globalActivity, index) => ( - - {globalActivity.model} - - - - API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} - - - - - - Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} - - - - - - ))} - - - - - - - - - - - Total Spend Per Team - - - - Daily Spend Per Team + +
+
+ + + Total Spend Per Team + + + + + + + + Daily Spend Per Team + + = ({ accessToken, token, userRole, use yAxisWidth={80} stack={true} /> - - - - - - -

- Customers of your LLM API calls. Tracked when a `user` param is passed in your LLM calls{" "} - - docs here - -

- - - { - setDateValue(value); - updateEndUserData(value.from, value.to, null); - }} - /> - - - Select Key - - - + + +
+
+ - - - - - Customer - Spend - Total Events - - - - - {topUsers?.map((user: any, index: number) => ( - - {user.end_user} - - - - {user.total_count} - + +

+ Customers of your LLM API calls. Tracked when a `user` param is passed in your LLM calls{" "} + + docs here + +

+
+
+ { + setDateValue(value); + updateEndUserData(value.from, value.to, null); + }} + /> +
+
+

Select Key

+
- - - - - - { - setDateValue(value); - updateTagSpendData(value.from, value.to); - }} - /> - + + +
+
- - {premiumUser ? ( -
- setSelectedTags(value as string[])}> - setSelectedTags(["all-tags"])} - > - All Tags - - {allTagNames && - allTagNames - .filter((tag) => tag !== "all-tags") - .map((tag: any, index: number) => { - return ( - - {tag} - - ); - })} - -
- ) : ( -
- setSelectedTags(value as string[])}> - setSelectedTags(["all-tags"])} - > - All Tags - - {allTagNames && - allTagNames - .filter((tag) => tag !== "all-tags") - .map((tag: any, index: number) => { - return ( - - ✨ {tag} (Enterprise only Feature) - - ); - })} - -
- )} - - - - - - Spend Per Tag - + + +
+ + + + Customer + Spend + Total Events + + + + + {topUsers?.map((user: any, index: number) => ( + + {user.end_user} + + + + {user.total_count} + + ))} + +
+
+ + + + + +
+
+ { + setDateValue(value); + updateTagSpendData(value.from, value.to); + }} + /> +
+ +
+ selectedTags.includes(option.value))} + onValueChange={(options: TagOption[]) => setSelectedTags(options.map((option) => option.value))} + isItemEqualToValue={(a: TagOption, b: TagOption) => a.value === b.value} + itemToStringLabel={(option: TagOption) => option.label} + > + + + {(options: TagOption[]) => + options.map((option) => ( + + {option.label} + + )) + } + + + + + No tags found + + {(option: TagOption) => ( + + {option.label} + + )} + + + +
+
+
+
+ + + Spend Per Tag + + +

Get Started by Tracking cost per tag{" "} here - - - - - - - - - +

+ + + +
+
+ +
); }; From 169ba0e287e9993f0b2f1a5d229c26d3d64d39bb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 22 Jul 2026 17:01:41 -0700 Subject: [PATCH 054/130] refactor(ui): migrate transform-request to shadcn (#34303) * test(ui): characterise transform-request panel behaviour before migration * refactor(ui): migrate transform-request to shadcn * fix(ui): keep transform-request panels within the fixed-height content fold * fix(ui): let transform-request flow naturally so the shell scrolls instead of clipping * test(ui): select the copy button by its accessible name --- ui/litellm-dashboard/eslint-suppressions.json | 5 - .../TransformRequestPanel.test.tsx | 160 ++++++++++++++ .../TransformRequestPanel.tsx | 205 ++++++------------ 3 files changed, 226 insertions(+), 144 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.test.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index b9b27e6cbb5..269348ec054 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1098,11 +1098,6 @@ "count": 1 } }, - "src/app/(dashboard)/transform-request/TransformRequestPanel.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/ui-theme/UIThemeSettings.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.test.tsx new file mode 100644 index 00000000000..a0add153116 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.test.tsx @@ -0,0 +1,160 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import TransformRequestPanel from "./TransformRequestPanel"; +import { transformRequestCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; + +vi.mock("@/components/networking", () => ({ + transformRequestCall: vi.fn(), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + info: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +const transformRequestCallMock = vi.mocked(transformRequestCall); +const notify = vi.mocked(NotificationsManager); + +const ACCESS_TOKEN = "sk-test-token"; + +const getRequestTextarea = () => screen.getByPlaceholderText(/press cmd\/ctrl \+ enter to transform/i); + +const getTransformButton = () => screen.getByRole("button", { name: /transform/i }); + +const getCopyButton = () => screen.getByRole("button", { name: /copy to clipboard/i }); + +describe("TransformRequestPanel", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("renders both panels, the prefilled request and the placeholder curl", () => { + render(); + + expect(screen.getByText("Original Request")).toBeInTheDocument(); + expect(screen.getByText("Transformed Request")).toBeInTheDocument(); + expect(screen.getByText(/sensitive headers are not shown/i)).toBeInTheDocument(); + + expect((getRequestTextarea() as HTMLTextAreaElement).value).toContain('"model": "openai/gpt-4o"'); + expect(screen.getByText(/https:\/\/api\.openai\.com\/v1\/chat\/completions/)).toBeInTheDocument(); + + expect(screen.getByRole("link", { name: /here/i })).toHaveAttribute( + "href", + "https://github.com/BerriAI/litellm/issues", + ); + }); + + it("sends the edited request body as a completion call and renders the returned curl", async () => { + const user = userEvent.setup(); + transformRequestCallMock.mockResolvedValue({ + raw_request_api_base: "https://api.anthropic.com/v1/messages", + raw_request_body: { model: "claude-opus-4-8", max_tokens: 42 }, + raw_request_headers: { "x-api-key": "redacted" }, + }); + + render(); + + const textarea = getRequestTextarea(); + await user.clear(textarea); + await user.type(textarea, '{{"model": "claude-opus-4-8"}'); + + await user.click(getTransformButton()); + + await waitFor(() => expect(transformRequestCallMock).toHaveBeenCalledTimes(1)); + expect(transformRequestCallMock).toHaveBeenCalledWith(ACCESS_TOKEN, { + call_type: "completion", + request_body: { model: "claude-opus-4-8" }, + }); + + const output = await screen.findByText(/api\.anthropic\.com\/v1\/messages/); + expect(output.textContent).toContain("curl -X POST"); + expect(output.textContent).toContain("-H 'x-api-key: redacted'"); + expect(output.textContent).toContain('"model": "claude-opus-4-8"'); + expect(output.textContent).toContain('"max_tokens": 42'); + expect(notify.success).toHaveBeenCalledWith("Request transformed successfully"); + }); + + it("transforms on Cmd/Ctrl + Enter without clicking the button", async () => { + const user = userEvent.setup(); + transformRequestCallMock.mockResolvedValue({ + raw_request_api_base: "https://api.openai.com/v1/chat/completions", + raw_request_body: { model: "gpt-4o" }, + raw_request_headers: {}, + }); + + render(); + + getRequestTextarea().focus(); + await user.keyboard("{Meta>}{Enter}{/Meta}"); + + await waitFor(() => expect(transformRequestCallMock).toHaveBeenCalledTimes(1)); + }); + + it("rejects invalid JSON without calling the backend", async () => { + const user = userEvent.setup(); + + render(); + + const textarea = getRequestTextarea(); + await user.clear(textarea); + await user.type(textarea, "not json"); + await user.click(getTransformButton()); + + await waitFor(() => expect(notify.fromBackend).toHaveBeenCalledWith("Invalid JSON in request body")); + expect(transformRequestCallMock).not.toHaveBeenCalled(); + }); + + it("does not call the backend when there is no access token", async () => { + const user = userEvent.setup(); + + render(); + + await user.click(getTransformButton()); + + await waitFor(() => expect(notify.fromBackend).toHaveBeenCalledWith("No access token found")); + expect(transformRequestCallMock).not.toHaveBeenCalled(); + }); + + it("reports a failed transform and leaves the placeholder curl in place", async () => { + const user = userEvent.setup(); + vi.spyOn(console, "error").mockImplementation(() => {}); + transformRequestCallMock.mockRejectedValue(new Error("boom")); + + render(); + + await user.click(getTransformButton()); + + await waitFor(() => expect(notify.fromBackend).toHaveBeenCalledWith("Failed to transform request")); + expect(screen.getByText(/https:\/\/api\.openai\.com\/v1\/chat\/completions/)).toBeInTheDocument(); + }); + + it("copies the transformed request to the clipboard", async () => { + const user = userEvent.setup(); + const writeText = vi.spyOn(navigator.clipboard, "writeText"); + transformRequestCallMock.mockResolvedValue({ + raw_request_api_base: "https://api.anthropic.com/v1/messages", + raw_request_body: { model: "claude-opus-4-8" }, + raw_request_headers: {}, + }); + + render(); + + await user.click(getTransformButton()); + await screen.findByText(/api\.anthropic\.com\/v1\/messages/); + + await user.click(getCopyButton()); + + expect(writeText).toHaveBeenCalledTimes(1); + expect(writeText.mock.calls[0]?.[0]).toContain("https://api.anthropic.com/v1/messages"); + expect(notify.success).toHaveBeenCalledWith("Copied to clipboard"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx index 04d1701de3f..0c41547b9b7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx @@ -1,9 +1,12 @@ import React, { useState } from "react"; -import { Button } from "antd"; -import { CopyOutlined } from "@ant-design/icons"; -import { Title } from "@tremor/react"; +import { ArrowRight, Copy } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardFooter, CardHeader, CardTitle } from "@/components/ui/card"; +import { Textarea } from "@/components/ui/textarea"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { transformRequestCall } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; + interface TransformRequestPanelProps { accessToken: string | null; } @@ -128,130 +131,50 @@ ${formattedBody} }; return ( -
- Playground -

See how LiteLLM transforms your request for the specified provider.

-
+
+

Playground

+

+ See how LiteLLM transforms your request for the specified provider. +

+
{/* Original Request Panel */} -
-
-

Original Request

-

- The request you would send to LiteLLM /chat/completions endpoint. -

-
+ + + Original Request + The request you would send to LiteLLM /chat/completions endpoint. + -